include/ck_tile/ops/gemm/warp/warp_gemm.hpp Source File

• include
• ck_tile
• ops
• gemm
• warp

warp_gemm.hpp

Go to the documentation of this file.

1 // SPDX-License-Identifier: MIT

2 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.

3 

4 #pragma once

5 

6 #include "ck_tile/core.hpp"

7 #include "ck_tile/ops/gemm/warp/warp_gemm_impl.hpp"

8 #include "ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp"

9 

10 namespace ck_tile {

11 

12 // fp16

13 

14 using WarpGemmMfmaF16F16F32M32N32K8 = WarpGemmImpl<

15  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImplF16F16F32M32N32K8<WGAttrCtlEnum::Default_>>>;

16 

17 using WarpGemmMfmaF16F16F32M16N16K16 = WarpGemmImpl<

18  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImplF16F16F32M16N16K16<WGAttrCtlEnum::Default_>>>;

19 

20 using WarpGemmMfmaF16F16F32M32N32K16 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

21  WarpGemmAttributeMfmaImplF16F16F32M32N32K8<WGAttrCtlEnum::Default_>,

22  2>>;

23 

24 using WarpGemmMfmaF16F16F32M16N16K32 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

25  WarpGemmAttributeMfmaImplF16F16F32M16N16K16<WGAttrCtlEnum::Default_>,

26  2>>;

27 

28 using WarpGemmMfmaF16F16F32M32N32K8SwizzleA = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK_SwizzleA<

29  WarpGemmAttributeMfmaImplF16F16F32M32N32K8<WGAttrCtlEnum::Default_>,

30  1>>;

31 

32 using WarpGemmMfmaF16F16F32M32N32K16SwizzleA = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK_SwizzleA<

33  WarpGemmAttributeMfmaImplF16F16F32M32N32K8<WGAttrCtlEnum::Default_>,

34  2>>;

35 

36 using WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution =

37  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

38  WarpGemmAttributeMfmaImplF16F16F32M32N32K8<WGAttrCtlEnum::Default_>>>;

39 

40 using WarpGemmMfmaF16F16F32M16N16K16TransposedCDistribution =

41  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

42  WarpGemmAttributeMfmaImplF16F16F32M16N16K16<WGAttrCtlEnum::Default_>>>;

43 

44 using WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution =

45  WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution<

46  WarpGemmAttributeMfmaImplF16F16F32M32N32K8<WGAttrCtlEnum::Default_>,

47  2>>;

48 

49 using WarpGemmMfmaF16F16F32M16N16K32TransposedCDistribution =

50  WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution<

51  WarpGemmAttributeMfmaImplF16F16F32M16N16K16<WGAttrCtlEnum::Default_>,

52  2>>;

53 

54 using WarpGemmMfmaF16F16F32M32N32K16SwizzleBTransposedCDistribution =

55  WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB<

56  WarpGemmAttributeMfmaImplF16F16F32M32N32K8<WGAttrCtlEnum::Default_>,

57  2>>;

58 

59 using WarpGemmMfmaF16F16F32M4N64K16 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

60  WarpGemmAttributeMfmaImplF16F16F32M4N64K4<WGAttrCtlEnum::Default_>,

61  4>>;

62 

63 using WarpGemmMfmaF16F16F32M64N4K16 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

64  WarpGemmAttributeMfmaImplF16F16F32M64N4K4<WGAttrCtlEnum::Default_>,

65  4>>;

66 

67 // bf16

68 

69 using WarpGemmMfmaBf16Bf16F32M32N32K8 = WarpGemmImpl<

70  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<WGAttrCtlEnum::Default_>>>;

71 

72 using WarpGemmMfmaBf16Bf16F32M16N16K16 = WarpGemmImpl<

73  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16<WGAttrCtlEnum::Default_>>>;

74 

75 using WarpGemmMfmaBf16Bf16F32M32N32K16 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

76  WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<WGAttrCtlEnum::Default_>,

77  2>>;

78 

79 using WarpGemmMfmaBf16Bf16F32M16N16K32 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

80  WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16<WGAttrCtlEnum::Default_>,

81  2>>;

82 

83 using WarpGemmMfmaBf16Bf16F32M32N32K8SwizzleA = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK_SwizzleA<

84  WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<WGAttrCtlEnum::Default_>,

85  1>>;

86 

87 using WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA =

88  WarpGemmImpl<WarpGemmAtrributeMfmaIterateK_SwizzleA<

89  WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<WGAttrCtlEnum::Default_>,

90  2>>;

91 

92 using WarpGemmMfmaBf16Bf16F32M32N32K8TransposedCDistribution =

93  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

94  WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<WGAttrCtlEnum::Default_>>>;

95 

96 using WarpGemmMfmaBf16Bf16F32M16N16K16TransposedCDistribution =

97  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

98  WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16<WGAttrCtlEnum::Default_>>>;

99 

100 using WarpGemmMfmaBf16Bf16F32M32N32K16TransposedCDistribution =

101  WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution<

102  WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<WGAttrCtlEnum::Default_>,

103  2>>;

104 

105 using WarpGemmMfmaBf16Bf16F32M16N16K32TransposedCDistribution =

106  WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution<

107  WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16<WGAttrCtlEnum::Default_>,

108  2>>;

109 

110 using WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleBTransposedCDistribution =

111  WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB<

112  WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<WGAttrCtlEnum::Default_>,

113  2>>;

114 

115 using WarpGemmMfmaBf16Bf16F32M4N64K16 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

116  WarpGemmAttributeMfmaImplBf16Bf16F32M4N64K4<WGAttrCtlEnum::Default_>,

117  4>>;

118 

119 using WarpGemmMfmaBf16Bf16F32M64N4K16 = WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<

120  WarpGemmAttributeMfmaImplBf16Bf16F32M64N4K4<WGAttrCtlEnum::Default_>,

121  4>>;

122 

123 // fp8

124 

125 using WarpGemmMfma_f32_32x32x16_fp8_fp8 = WarpGemmImpl<

126  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_fp8<WGAttrCtlEnum::Default_>>>;

127 

128 using WarpGemmMfma_f32_32x32x16_fp8_bf8 = WarpGemmImpl<

129  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_bf8<WGAttrCtlEnum::Default_>>>;

130 

131 using WarpGemmMfma_f32_32x32x16_bf8_fp8 = WarpGemmImpl<

132  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImpl_f32_32x32x16_bf8_fp8<WGAttrCtlEnum::Default_>>>;

133 

134 using WarpGemmMfma_f32_32x32x16_bf8_bf8 = WarpGemmImpl<

135  WarpGemmAtrributeMfma<WarpGemmAttributeMfmaImpl_f32_32x32x16_bf8_bf8<WGAttrCtlEnum::Default_>>>;

136 

137 using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed =

138  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

139  WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_fp8<WGAttrCtlEnum::Default_>>>;

140 

141 using WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed =

142  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

143  WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_bf8<WGAttrCtlEnum::Default_>>>;

144 

145 using WarpGemmMfma_f32_32x32x16_bf8_fp8_CTransposed =

146  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

147  WarpGemmAttributeMfmaImpl_f32_32x32x16_bf8_fp8<WGAttrCtlEnum::Default_>>>;

148 

149 using WarpGemmMfma_f32_32x32x16_bf8_bf8_CTransposed =

150  WarpGemmImpl<WarpGemmAtrributeMfmaTransposedCDistribution<

151  WarpGemmAttributeMfmaImpl_f32_32x32x16_bf8_bf8<WGAttrCtlEnum::Default_>>>;

152 

153 template <index_t swizzle_factor = 2>

154 using WarpGemmMfmaFp8Fp8F32M32N32K16SwizzleBTransposedCDistribution =

155  WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB<

156  WarpGemmAttributeMfmaImpl_f32_32x32x16_f8_base<fp8_t, fp8_t, WGAttrCtlEnum::Default_>,

157  2,

158  swizzle_factor>>;

159 

160 } // namespace ck_tile

core.hpp

ck_tile

Definition: cluster_descriptor.hpp:13

ck_tile::WarpGemmAtrributeMfma

Definition: warp_gemm_attribute_mfma.hpp:13

ck_tile::WarpGemmAtrributeMfmaIterateK_SwizzleA

Definition: warp_gemm_attribute_mfma.hpp:791

ck_tile::WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB

Definition: warp_gemm_attribute_mfma.hpp:647

ck_tile::WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution

Definition: warp_gemm_attribute_mfma.hpp:432

ck_tile::WarpGemmAtrributeMfmaIterateK

Definition: warp_gemm_attribute_mfma.hpp:78

ck_tile::WarpGemmAtrributeMfmaTransposedCDistribution

Definition: warp_gemm_attribute_mfma.hpp:294

ck_tile::WarpGemmAttributeMfmaImpl_f32_32x32x16_f8_base

Definition: warp_gemm_attribute_mfma_impl.hpp:628

ck_tile::WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16

Definition: warp_gemm_attribute_mfma_impl.hpp:411

ck_tile::WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8

Definition: warp_gemm_attribute_mfma_impl.hpp:322

ck_tile::WarpGemmAttributeMfmaImplBf16Bf16F32M4N64K4

Definition: warp_gemm_attribute_mfma_impl.hpp:499

ck_tile::WarpGemmAttributeMfmaImplBf16Bf16F32M64N4K4

Definition: warp_gemm_attribute_mfma_impl.hpp:563

ck_tile::WarpGemmAttributeMfmaImplF16F16F32M16N16K16

Definition: warp_gemm_attribute_mfma_impl.hpp:130

ck_tile::WarpGemmAttributeMfmaImplF16F16F32M32N32K8

Definition: warp_gemm_attribute_mfma_impl.hpp:67

ck_tile::WarpGemmAttributeMfmaImplF16F16F32M4N64K4

Definition: warp_gemm_attribute_mfma_impl.hpp:193

ck_tile::WarpGemmAttributeMfmaImplF16F16F32M64N4K4

Definition: warp_gemm_attribute_mfma_impl.hpp:257

ck_tile::WarpGemmImpl

Definition: warp_gemm_impl.hpp:11

warp_gemm_attribute_mfma.hpp

warp_gemm_impl.hpp