transform_conv_fwd_to_gemm.hpp Source File

transform_conv_fwd_to_gemm.hpp Source File#

Composable Kernel: transform_conv_fwd_to_gemm.hpp Source File
tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp
Go to the documentation of this file.
1
2// SPDX-License-Identifier: MIT
3// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
4
5#pragma once
6
7#include "ck/utility/common_header.hpp"
8#include "ck/tensor_description/tensor_descriptor.hpp"
9#include "ck/tensor_description/tensor_descriptor_helper.hpp"
10#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
11#include "ck/tensor_operation/gpu/device/convolution_forward_specialization.hpp"
12
13namespace ck {
14namespace tensor_operation {
15
16template <index_t NDimSpatial,
17 device::ConvolutionForwardSpecialization ConvForwardSpecialization,
18 bool SplitN = false,
19 typename ADataType = float,
20 typename CDataType = float,
21 index_t NumGroupsToMerge = 1,
22 typename IndexType = index_t,
23 bool CTranspose = false>
24struct TransformConvFwdToGemm
25{
26 private:
27 static constexpr auto I0 = Number<0>{};
28 static constexpr auto I1 = Number<1>{};
29 static constexpr auto I2 = Number<2>{};
30 static constexpr auto I3 = Number<3>{};
31 static constexpr auto I4 = Number<4>{};
32 static constexpr auto I5 = Number<5>{};
33
34 template <typename ConvDimsType>
35 static long_index_t calculate_element_space_size_impl(const ConvDimsType& lengths,
36 const ConvDimsType& strides,
37 index_t i)
38 {
39 long_index_t acc = 1;
40 for(; i < (NDimSpatial + 3); i++)
41 {
42 acc +=
43 static_cast<long_index_t>(lengths[i] - I1) * static_cast<long_index_t>(strides[i]);
44 }
45
46 return acc;
47 }
48
49 template <typename ConvDimsType>
50 static IndexType GetSplitedNSize(const ConvDimsType& a_g_n_c_wis_lengths,
51 const ConvDimsType& a_g_n_c_wis_strides,
52 const ConvDimsType& c_g_n_k_wos_lengths,
53 const ConvDimsType& c_g_n_k_wos_strides)
54 {
55 const long_index_t a_element_space_size =
56 calculate_element_space_size_impl(a_g_n_c_wis_lengths, a_g_n_c_wis_strides, I1);
57 const long_index_t c_element_space_size =
58 calculate_element_space_size_impl(c_g_n_k_wos_lengths, c_g_n_k_wos_strides, I1);
59 const long_index_t element_space_size = math::max(a_element_space_size * sizeof(ADataType),
60 c_element_space_size * sizeof(CDataType));
61 constexpr long_index_t TwoGB = (long_index_t{1} << 31); // 2GB threshold
62
63 const IndexType N = a_g_n_c_wis_lengths[I1];
64
65 if(element_space_size > TwoGB)
66 {
67 // Minimum divisor of N to not exceed 2GB
68 const auto divisor = math::integer_divide_ceil(element_space_size, TwoGB);
69
70 if(divisor <= static_cast<double>(N))
71 {
72 // Find least divisor of N larger than element_space_size / TwoGB
73 // Iterate up to sqrt(N). There are no divisors above this value.
74 for(IndexType least_divisor = divisor; least_divisor * least_divisor <= N;
75 least_divisor++)
76 {
77 if(N % least_divisor == 0)
78 {
79 return N / least_divisor;
80 }
81 }
82 // Not found, process one Convolution N per block
83 return 1;
84 }
85 else
86 {
87 // Split Convolution's N dimension into N workgroups. However
88 // this still might not result in sufficiently small tensor,
89 // but at least later on we could divide the image as well.
90 return 1;
91 }
92 }
93 else
94 {
95 // Split N is not needed.
96 return N;
97 }
98 }
99
100 public:
101 __host__ __device__ constexpr TransformConvFwdToGemm() {}
102
103 template <typename TransformConvFwdToGemmBase>
104 __host__ __device__
105 TransformConvFwdToGemm(const TransformConvFwdToGemmBase& transform_conv_fwd_to_gemm_base)
106 : N_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.N_)},
107 Di_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Di_)},
108 Hi_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Hi_)},
109 Wi_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Wi_)},
110 Do_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Do_)},
111 Ho_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Ho_)},
112 Wo_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Wo_)},
113 Z_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Z_)},
114 Y_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.Y_)},
115 X_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.X_)},
116 K_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.K_)},
117 C_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.C_)},
118 DiStride_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.DiStride_)},
119 HiStride_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.HiStride_)},
120 WiStride_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.WiStride_)},
121 DoStride_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.DoStride_)},
122 HoStride_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.HoStride_)},
123 WoStride_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.WoStride_)},
124 XStride_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.XStride_)},
125 CStrideTensorA_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.CStrideTensorA_)},
126 CStrideTensorB_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.CStrideTensorB_)},
127 KStrideTensorB_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.KStrideTensorB_)},
128 KStrideTensorC_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.KStrideTensorC_)},
129 NStrideTensorA_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.NStrideTensorA_)},
130 NStrideTensorC_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.NStrideTensorC_)},
131 GStrideTensorA_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.GStrideTensorA_)},
132 GStrideTensorB_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.GStrideTensorB_)},
133 GStrideTensorC_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.GStrideTensorC_)},
134 ConvStrideD_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.ConvStrideD_)},
135 ConvStrideH_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.ConvStrideH_)},
136 ConvStrideW_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.ConvStrideW_)},
137 ConvDilationD_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.ConvDilationD_)},
138 ConvDilationH_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.ConvDilationH_)},
139 ConvDilationW_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.ConvDilationW_)},
140 InLeftPadD_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.InLeftPadD_)},
141 InLeftPadH_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.InLeftPadH_)},
142 InLeftPadW_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.InLeftPadW_)},
143 InRightPadD_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.InRightPadD_)},
144 InRightPadH_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.InRightPadH_)},
145 InRightPadW_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.InRightPadW_)},
146 ZYX_{static_cast<IndexType>(transform_conv_fwd_to_gemm_base.ZYX_)}
147 {
148 }
149
150 template <typename ConvDimsType,
151 typename ConvSpatialDimsType,
152 index_t NDim = NDimSpatial,
153 typename ck::enable_if<NDim == 1, bool>::type = false>
154 __host__ __device__ TransformConvFwdToGemm(const ConvDimsType& a_g_n_c_wis_lengths,
155 const ConvDimsType& a_g_n_c_wis_strides,
156 const ConvDimsType& b_g_k_c_xs_lengths,
157 const ConvDimsType& b_g_k_c_xs_strides,
158 const ConvDimsType& c_g_n_k_wos_lengths,
159 const ConvDimsType& c_g_n_k_wos_strides,
160 const ConvSpatialDimsType& conv_filter_strides,
161 const ConvSpatialDimsType& conv_filter_dilations,
162 const ConvSpatialDimsType& input_left_pads,
163 const ConvSpatialDimsType& input_right_pads)
164 : Di_{I1},
165 Hi_{I1},
166 Wi_{a_g_n_c_wis_lengths[I3]},
167 Do_{I1},
168 Ho_{I1},
169 Wo_{c_g_n_k_wos_lengths[I3]},
170 Z_{I1},
171 Y_{I1},
172 X_{b_g_k_c_xs_lengths[I3]},
173 K_{c_g_n_k_wos_lengths[I2]},
174 C_{b_g_k_c_xs_lengths[I2]},
175 DiStride_{I1},
176 HiStride_{I1},
177 WiStride_{a_g_n_c_wis_strides[I3]},
178 DoStride_{I1},
179 HoStride_{I1},
180 WoStride_{c_g_n_k_wos_strides[I3]},
181 XStride_{b_g_k_c_xs_strides[I3]},
182 CStrideTensorA_{a_g_n_c_wis_strides[I2]},
183 CStrideTensorB_{b_g_k_c_xs_strides[I2]},
184 KStrideTensorB_{b_g_k_c_xs_strides[I1]},
185 KStrideTensorC_{c_g_n_k_wos_strides[I2]},
186 NStrideTensorA_{a_g_n_c_wis_strides[I1]},
187 NStrideTensorC_{c_g_n_k_wos_strides[I1]},
188 GStrideTensorA_{a_g_n_c_wis_strides[I0]},
189 GStrideTensorB_{b_g_k_c_xs_strides[I0]},
190 GStrideTensorC_{c_g_n_k_wos_strides[I0]},
191 ConvStrideD_{I1},
192 ConvStrideH_{I1},
193 ConvStrideW_{conv_filter_strides[I0]},
194 ConvDilationD_{I1},
195 ConvDilationH_{I1},
196 ConvDilationW_{conv_filter_dilations[I0]},
197 InLeftPadD_{I0},
198 InLeftPadH_{I0},
199 InLeftPadW_{input_left_pads[I0]},
200 InRightPadD_{I0},
201 InRightPadH_{I0},
202 InRightPadW_{input_right_pads[I0]},
203 ZYX_{X_}
204 {
205#ifdef CK_CODE_GEN_RTC
206 static_assert(is_same_v<ConvSpatialDimsType, ck::Array<IndexType, NDimSpatial>>);
207 static_assert(is_same_v<ConvDimsType, ck::Array<IndexType, NDimSpatial + I3>>);
208#else
209 static_assert(is_same_v<ConvSpatialDimsType, std::array<IndexType, NDimSpatial>> ||
210 is_same_v<ConvSpatialDimsType, ck::Array<IndexType, NDimSpatial>>);
211 static_assert(is_same_v<ConvDimsType, std::array<IndexType, NDimSpatial + I3>> ||
212 is_same_v<ConvDimsType, ck::Array<IndexType, NDimSpatial + I3>>);
213#endif
214 if constexpr(SplitN)
215 {
216 N_ = GetSplitedNSize(
217 a_g_n_c_wis_lengths, a_g_n_c_wis_strides, c_g_n_k_wos_lengths, c_g_n_k_wos_strides);
218 }
219 else
220 {
221 N_ = c_g_n_k_wos_lengths[I1];
222 }
223 }
224
225 template <typename ConvDimsType,
226 typename ConvSpatialDimsType,
227 index_t NDim = NDimSpatial,
228 typename ck::enable_if<NDim == 2, bool>::type = false>
229 __host__ __device__ TransformConvFwdToGemm(const ConvDimsType& a_g_n_c_wis_lengths,
230 const ConvDimsType& a_g_n_c_wis_strides,
231 const ConvDimsType& b_g_k_c_xs_lengths,
232 const ConvDimsType& b_g_k_c_xs_strides,
233 const ConvDimsType& c_g_n_k_wos_lengths,
234 const ConvDimsType& c_g_n_k_wos_strides,
235 const ConvSpatialDimsType& conv_filter_strides,
236 const ConvSpatialDimsType& conv_filter_dilations,
237 const ConvSpatialDimsType& input_left_pads,
238 const ConvSpatialDimsType& input_right_pads)
239 : Di_{I1},
240 Hi_{a_g_n_c_wis_lengths[I3]},
241 Wi_{a_g_n_c_wis_lengths[I4]},
242 Do_{I1},
243 Ho_{c_g_n_k_wos_lengths[I3]},
244 Wo_{c_g_n_k_wos_lengths[I4]},
245 Z_{I1},
246 Y_{b_g_k_c_xs_lengths[I3]},
247 X_{b_g_k_c_xs_lengths[I4]},
248 K_{c_g_n_k_wos_lengths[I2]},
249 C_{b_g_k_c_xs_lengths[I2]},
250 DiStride_{I1},
251 HiStride_{a_g_n_c_wis_strides[I3]},
252 WiStride_{a_g_n_c_wis_strides[I4]},
253 DoStride_{I1},
254 HoStride_{c_g_n_k_wos_strides[I3]},
255 WoStride_{c_g_n_k_wos_strides[I4]},
256 XStride_{b_g_k_c_xs_strides[I4]},
257 CStrideTensorA_{a_g_n_c_wis_strides[I2]},
258 CStrideTensorB_{b_g_k_c_xs_strides[I2]},
259 KStrideTensorB_{b_g_k_c_xs_strides[I1]},
260 KStrideTensorC_{c_g_n_k_wos_strides[I2]},
261 NStrideTensorA_{a_g_n_c_wis_strides[I1]},
262 NStrideTensorC_{c_g_n_k_wos_strides[I1]},
263 GStrideTensorA_{a_g_n_c_wis_strides[I0]},
264 GStrideTensorB_{b_g_k_c_xs_strides[I0]},
265 GStrideTensorC_{c_g_n_k_wos_strides[I0]},
266 ConvStrideD_{I1},
267 ConvStrideH_{conv_filter_strides[I0]},
268 ConvStrideW_{conv_filter_strides[I1]},
269 ConvDilationD_{I1},
270 ConvDilationH_{conv_filter_dilations[I0]},
271 ConvDilationW_{conv_filter_dilations[I1]},
272 InLeftPadD_{I0},
273 InLeftPadH_{input_left_pads[I0]},
274 InLeftPadW_{input_left_pads[I1]},
275 InRightPadD_{I0},
276 InRightPadH_{input_right_pads[I0]},
277 InRightPadW_{input_right_pads[I1]},
278 ZYX_{Y_ * X_}
279 {
280#ifdef CK_CODE_GEN_RTC
281 static_assert(is_same_v<ConvSpatialDimsType, ck::Array<IndexType, NDimSpatial>>);
282 static_assert(is_same_v<ConvDimsType, ck::Array<IndexType, NDimSpatial + I3>>);
283#else
284 static_assert(is_same_v<ConvSpatialDimsType, std::array<IndexType, NDimSpatial>> ||
285 is_same_v<ConvSpatialDimsType, ck::Array<IndexType, NDimSpatial>>);
286 static_assert(is_same_v<ConvDimsType, std::array<IndexType, NDimSpatial + I3>> ||
287 is_same_v<ConvDimsType, ck::Array<IndexType, NDimSpatial + I3>>);
288#endif
289 if constexpr(SplitN)
290 {
291 N_ = GetSplitedNSize(
292 a_g_n_c_wis_lengths, a_g_n_c_wis_strides, c_g_n_k_wos_lengths, c_g_n_k_wos_strides);
293 }
294 else
295 {
296 N_ = c_g_n_k_wos_lengths[I1];
297 }
298 }
299
300 template <typename ConvDimsType,
301 typename ConvSpatialDimsType,
302 index_t NDim = NDimSpatial,
303 typename ck::enable_if<NDim == 3, bool>::type = false>
304 __host__ __device__ TransformConvFwdToGemm(const ConvDimsType& a_g_n_c_wis_lengths,
305 const ConvDimsType& a_g_n_c_wis_strides,
306 const ConvDimsType& b_g_k_c_xs_lengths,
307 const ConvDimsType& b_g_k_c_xs_strides,
308 const ConvDimsType& c_g_n_k_wos_lengths,
309 const ConvDimsType& c_g_n_k_wos_strides,
310 const ConvSpatialDimsType& conv_filter_strides,
311 const ConvSpatialDimsType& conv_filter_dilations,
312 const ConvSpatialDimsType& input_left_pads,
313 const ConvSpatialDimsType& input_right_pads)
314 : Di_{a_g_n_c_wis_lengths[I3]},
315 Hi_{a_g_n_c_wis_lengths[I4]},
316 Wi_{a_g_n_c_wis_lengths[I5]},
317 Do_{c_g_n_k_wos_lengths[I3]},
318 Ho_{c_g_n_k_wos_lengths[I4]},
319 Wo_{c_g_n_k_wos_lengths[I5]},
320 Z_{b_g_k_c_xs_lengths[I3]},
321 Y_{b_g_k_c_xs_lengths[I4]},
322 X_{b_g_k_c_xs_lengths[I5]},
323 K_{c_g_n_k_wos_lengths[I2]},
324 C_{b_g_k_c_xs_lengths[I2]},
325 DiStride_{a_g_n_c_wis_strides[I3]},
326 HiStride_{a_g_n_c_wis_strides[I4]},
327 WiStride_{a_g_n_c_wis_strides[I5]},
328 DoStride_{c_g_n_k_wos_strides[I3]},
329 HoStride_{c_g_n_k_wos_strides[I4]},
330 WoStride_{c_g_n_k_wos_strides[I5]},
331 XStride_{b_g_k_c_xs_strides[I5]},
332 CStrideTensorA_{a_g_n_c_wis_strides[I2]},
333 CStrideTensorB_{b_g_k_c_xs_strides[I2]},
334 KStrideTensorB_{b_g_k_c_xs_strides[I1]},
335 KStrideTensorC_{c_g_n_k_wos_strides[I2]},
336 NStrideTensorA_{a_g_n_c_wis_strides[I1]},
337 NStrideTensorC_{c_g_n_k_wos_strides[I1]},
338 GStrideTensorA_{a_g_n_c_wis_strides[I0]},
339 GStrideTensorB_{b_g_k_c_xs_strides[I0]},
340 GStrideTensorC_{c_g_n_k_wos_strides[I0]},
341 ConvStrideD_{conv_filter_strides[I0]},
342 ConvStrideH_{conv_filter_strides[I1]},
343 ConvStrideW_{conv_filter_strides[I2]},
344 ConvDilationD_{conv_filter_dilations[I0]},
345 ConvDilationH_{conv_filter_dilations[I1]},
346 ConvDilationW_{conv_filter_dilations[I2]},
347 InLeftPadD_{input_left_pads[I0]},
348 InLeftPadH_{input_left_pads[I1]},
349 InLeftPadW_{input_left_pads[I2]},
350 InRightPadD_{input_right_pads[I0]},
351 InRightPadH_{input_right_pads[I1]},
352 InRightPadW_{input_right_pads[I2]},
353 ZYX_{Z_ * Y_ * X_}
354 {
355#ifdef CK_CODE_GEN_RTC
356 static_assert(is_same_v<ConvSpatialDimsType, ck::Array<IndexType, NDimSpatial>>);
357 static_assert(is_same_v<ConvDimsType, ck::Array<IndexType, NDimSpatial + I3>>);
358#else
359 static_assert(is_same_v<ConvSpatialDimsType, std::array<IndexType, NDimSpatial>> ||
360 is_same_v<ConvSpatialDimsType, ck::Array<IndexType, NDimSpatial>>);
361 static_assert(is_same_v<ConvDimsType, std::array<IndexType, NDimSpatial + I3>> ||
362 is_same_v<ConvDimsType, ck::Array<IndexType, NDimSpatial + I3>>);
363#endif
364 if constexpr(SplitN)
365 {
366 N_ = GetSplitedNSize(
367 a_g_n_c_wis_lengths, a_g_n_c_wis_strides, c_g_n_k_wos_lengths, c_g_n_k_wos_strides);
368 }
369 else
370 {
371 N_ = c_g_n_k_wos_lengths[I1];
372 }
373 }
374
375 __host__ bool AreDescriptorsSmallerThan2GB() const
376 {
377 constexpr long_index_t TwoGB = (long_index_t{1} << 31);
378
379 const long_index_t in_desc_space_size =
380 I1 + (N_ - I1) * NStrideTensorA_ + (Di_ - I1) * DiStride_ + (Hi_ - I1) * HiStride_ +
381 (Wi_ - I1) * WiStride_ + (C_ - I1) * CStrideTensorA_;
382 const long_index_t out_desc_space_size =
383 I1 + (N_ - I1) * NStrideTensorC_ + (Do_ - I1) * DoStride_ + (Ho_ - I1) * HoStride_ +
384 (Wo_ - I1) * WoStride_ + (K_ - I1) * KStrideTensorC_;
385
386 bool is_a_descriptor_smaller_than_2GB = (in_desc_space_size * sizeof(ADataType)) <= TwoGB;
387 bool is_c_descriptor_smaller_than_2GB = (out_desc_space_size * sizeof(CDataType)) <= TwoGB;
388
389 return is_a_descriptor_smaller_than_2GB && is_c_descriptor_smaller_than_2GB;
390 }
391
392 template <typename DsPointer>
393 __host__ auto SplitConvProblem(const ADataType* a_grid_ptr_base,
394 DsPointer& ds_grid_ptr_base,
395 CDataType* c_grid_ptr_base) const
396 {
397 // Create copies
398 auto conv_to_gemm_transformer_left = *this;
399 auto conv_to_gemm_transformer_right = *this;
400 IndexType a_right_offset = 0;
401 IndexType c_right_offset = 0;
402 // Calculate real filter size
403 const IndexType z_eff = (Z_ - 1) * ConvDilationD_ + 1;
404 const IndexType y_eff = (Y_ - 1) * ConvDilationH_ + 1;
405 const IndexType x_eff = (X_ - 1) * ConvDilationW_ + 1;
406 // Calculate start position in input for right tensor
407 const IndexType di_right_transformer_start_idx = (Do_ / 2) * ConvStrideD_;
408 const IndexType hi_right_transformer_start_idx = (Ho_ / 2) * ConvStrideH_;
409 const IndexType wi_right_transformer_start_idx = (Wo_ / 2) * ConvStrideW_;
410 // Calculate last position in input for left tensor
411 const IndexType di_left_transformer_end_idx = (Do_ / 2 - 1) * ConvStrideD_ + z_eff;
412 const IndexType hi_left_transformer_end_idx = (Ho_ / 2 - 1) * ConvStrideH_ + y_eff;
413 const IndexType wi_left_transformer_end_idx = (Wo_ / 2 - 1) * ConvStrideW_ + x_eff;
414 // Allow to split if whole left padding will be in left tensor and right padding in right
415 // tensor
416 const bool is_possible_to_split_d = Do_ != 1 &&
417 di_right_transformer_start_idx > InLeftPadD_ &&
418 di_left_transformer_end_idx <= (InLeftPadD_ + Di_);
419 const bool is_possible_to_split_h = Ho_ != 1 &&
420 hi_right_transformer_start_idx > InLeftPadH_ &&
421 hi_left_transformer_end_idx <= (InLeftPadH_ + Hi_);
422 const bool is_possible_to_split_w = Wo_ != 1 &&
423 wi_right_transformer_start_idx > InLeftPadW_ &&
424 wi_left_transformer_end_idx <= (InLeftPadW_ + Wi_);
425
426 if(is_possible_to_split_d)
427 {
428 // Apply new sizes
429 // Split output on half
430 conv_to_gemm_transformer_left.Do_ = Do_ / 2;
431 conv_to_gemm_transformer_right.Do_ = Do_ - Do_ / 2;
432 // Assign left padding to left convolution
433 conv_to_gemm_transformer_left.InLeftPadD_ = InLeftPadD_;
434 conv_to_gemm_transformer_right.InLeftPadD_ = 0;
435 // Assign right padding to right convolution
436 conv_to_gemm_transformer_left.InRightPadD_ = 0;
437 conv_to_gemm_transformer_right.InRightPadD_ = InRightPadD_;
438 // Calculate new input size
439 conv_to_gemm_transformer_left.Di_ = di_left_transformer_end_idx - InLeftPadD_;
440 conv_to_gemm_transformer_right.Di_ =
441 math::min(Di_ - (di_right_transformer_start_idx - InLeftPadD_),
442 (conv_to_gemm_transformer_right.Do_ - 1) * ConvStrideD_ + z_eff);
443 ;
444 // Calcualte offsets
445 a_right_offset = ((Do_ / 2) * ConvStrideD_ - InLeftPadD_) * DiStride_;
446 c_right_offset = (Do_ / 2) * DoStride_;
447 }
448 else if(is_possible_to_split_h)
449 {
450 conv_to_gemm_transformer_left.Ho_ = Ho_ / 2;
451 conv_to_gemm_transformer_right.Ho_ = Ho_ - Ho_ / 2;
452
453 conv_to_gemm_transformer_left.InLeftPadH_ = InLeftPadH_;
454 conv_to_gemm_transformer_right.InLeftPadH_ = 0;
455
456 conv_to_gemm_transformer_left.InRightPadH_ = 0;
457 conv_to_gemm_transformer_right.InRightPadH_ = InRightPadH_;
458
459 conv_to_gemm_transformer_left.Hi_ = hi_left_transformer_end_idx - InLeftPadH_;
460 conv_to_gemm_transformer_right.Hi_ =
461 math::min(Hi_ - (hi_right_transformer_start_idx - InLeftPadH_),
462 (conv_to_gemm_transformer_right.Ho_ - 1) * ConvStrideH_ + y_eff);
463 a_right_offset = ((Ho_ / 2) * ConvStrideH_ - InLeftPadH_) * HiStride_;
464 c_right_offset = (Ho_ / 2) * HoStride_;
465 }
466 else if(is_possible_to_split_w)
467 {
468 conv_to_gemm_transformer_left.Wo_ = Wo_ / 2;
469 conv_to_gemm_transformer_right.Wo_ = Wo_ - Wo_ / 2;
470
471 conv_to_gemm_transformer_left.InLeftPadW_ = InLeftPadW_;
472 conv_to_gemm_transformer_right.InLeftPadW_ = 0;
473
474 conv_to_gemm_transformer_left.InRightPadW_ = 0;
475 conv_to_gemm_transformer_right.InRightPadW_ = InRightPadW_;
476
477 conv_to_gemm_transformer_left.Wi_ = wi_left_transformer_end_idx - InLeftPadW_;
478 conv_to_gemm_transformer_right.Wi_ =
479 math::min(Wi_ - (wi_right_transformer_start_idx - InLeftPadW_),
480 (conv_to_gemm_transformer_right.Wo_ - 1) * ConvStrideW_ + x_eff);
481
482 a_right_offset = ((Wo_ / 2) * ConvStrideW_ - InLeftPadW_) * WiStride_;
483 c_right_offset = (Wo_ / 2) * WoStride_;
484 }
485
486 static constexpr index_t NumDTensor = DsPointer::Size();
487 const auto ds_grid_right_ptr = generate_tuple(
488 [&](auto i) { return ds_grid_ptr_base(i) + c_right_offset; }, Number<NumDTensor>{});
489
490 // Return left transform, right transformer, right offset to Input and right offset to
491 // Output
492 return ck::make_tuple(conv_to_gemm_transformer_left,
493 conv_to_gemm_transformer_right,
494 a_grid_ptr_base + a_right_offset,
495 ds_grid_right_ptr,
496 c_grid_ptr_base + c_right_offset);
497 }
498
499 // TODO: implement ck::tensor_layout::convolution that describe packed/strided dimemsion as
500 // properties
501 template <typename ALayout,
502 typename ck::enable_if<NDimSpatial == 1 &&
503 (is_same_v<ALayout, tensor_layout::convolution::G_NW_C> ||
504 is_same_v<ALayout, tensor_layout::convolution::NWGC> ||
505 is_same_v<ALayout, tensor_layout::convolution::GNWC>),
506 bool>::type = false>
507 __host__ __device__ auto MakeADescriptor_M_K() const
508 {
509 if constexpr(ConvForwardSpecialization ==
510 device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
511 {
512 if constexpr(NumGroupsToMerge == 1)
513 {
514 const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor(
515 make_tuple(N_, Wo_, C_),
516 make_tuple(NStrideTensorA_, WiStride_, CStrideTensorA_));
517 return transform_tensor_descriptor(
518 in_gemmm_gemmk_desc,
519 make_tuple(make_merge_transform(make_tuple(N_, Wo_)),
520 make_pass_through_transform(C_)),
521 make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
522 make_tuple(Sequence<0>{}, Sequence<1>{}));
523 }
524 else
525 {
526 const auto in_gemmm_groups_gemmk_desc = make_naive_tensor_descriptor(
527 make_tuple(N_, Wo_, NumGroupsToMerge, C_),
528 make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_));
529
530 return transform_tensor_descriptor(
531 in_gemmm_groups_gemmk_desc,
532 make_tuple(make_merge_transform(make_tuple(N_, Wo_, NumGroupsToMerge)),
533 make_pass_through_transform(C_)),
534 make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}),
535 make_tuple(Sequence<0>{}, Sequence<1>{}));
536 }
537 }
538 else if constexpr(ConvForwardSpecialization ==
539 device::ConvolutionForwardSpecialization::Filter3x3)
540 {
541 if constexpr(NumGroupsToMerge == 1)
542 {
543
544 const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
545 make_tuple(N_, Wi_), make_tuple(NStrideTensorA_, WiStride_));
546
547 const auto in_n_wip_c_desc = transform_tensor_descriptor(
548 in_n_wi_c_desc,
549 make_tuple(make_pass_through_transform(N_),
550 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_)),
551 make_tuple(Sequence<0>{}, Sequence<1>{}),
552 make_tuple(Sequence<0>{}, Sequence<1>{}));
553
554 const auto in_n_x_wo_c_desc = transform_tensor_descriptor(
555 in_n_wip_c_desc,
556 make_tuple(make_pass_through_transform(N_),
557 make_embed_transform(make_tuple(Number<3>{}, Wo_),
558 make_tuple(ConvDilationW_, ConvStrideW_))),
559 make_tuple(Sequence<0>{}, Sequence<1>{}),
560 make_tuple(Sequence<0>{}, Sequence<1, 2>{}));
561
562 return transform_tensor_descriptor(
563 in_n_x_wo_c_desc,
564 make_tuple(make_merge_transform(make_tuple(N_, Wo_)),
565 make_pass_through_transform(Number<3>{})),
566 make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
567 make_tuple(Sequence<0>{}, Sequence<1>{}));
568 }
569 else
570 {
571 const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
572 make_tuple(N_, Wi_, NumGroupsToMerge),
573 make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_));
574
575 const auto in_n_wip_c_desc = transform_tensor_descriptor(
576 in_n_wi_c_desc,
577 make_tuple(make_pass_through_transform(N_),
578 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
579 make_pass_through_transform(NumGroupsToMerge)),
580 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
581 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
582
583 const auto in_n_x_wo_c_desc = transform_tensor_descriptor(
584 in_n_wip_c_desc,
585 make_tuple(make_pass_through_transform(N_),
586 make_embed_transform(make_tuple(Number<3>{}, Wo_),
587 make_tuple(ConvDilationW_, ConvStrideW_)),
588 make_pass_through_transform(NumGroupsToMerge)),
589 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
590 make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}));
591
592 return transform_tensor_descriptor(
593 in_n_x_wo_c_desc,
594 make_tuple(make_merge_transform(make_tuple(N_, Wo_, NumGroupsToMerge)),
595 make_pass_through_transform(Number<3>{})),
596 make_tuple(Sequence<0, 2, 3>{}, Sequence<1>{}),
597 make_tuple(Sequence<0>{}, Sequence<1>{}));
598 }
599 }
600 else if constexpr(ConvForwardSpecialization ==
601 device::ConvolutionForwardSpecialization::Filter1x1Pad0)
602 {
603 if constexpr(NumGroupsToMerge == 1)
604 {
605 const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
606 make_tuple(N_, Wi_, C_),
607 make_tuple(NStrideTensorA_, WiStride_, CStrideTensorA_));
608
609 const auto in_n_wo_c_desc = transform_tensor_descriptor(
610 in_n_wi_c_desc,
611 make_tuple(make_pass_through_transform(N_),
612 make_embed_transform(make_tuple(Wo_), make_tuple(ConvStrideW_)),
613 make_pass_through_transform(C_)),
614 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
615 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
616
617 return transform_tensor_descriptor(
618 in_n_wo_c_desc,
619 make_tuple(make_merge_transform(make_tuple(N_, Wo_)),
620 make_pass_through_transform(C_)),
621 make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
622 make_tuple(Sequence<0>{}, Sequence<1>{}));
623 }
624 else
625 {
626 const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
627 make_tuple(N_, Wi_, NumGroupsToMerge, C_),
628 make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_));
629
630 const auto in_n_wo_c_desc = transform_tensor_descriptor(
631 in_n_wi_c_desc,
632 make_tuple(make_pass_through_transform(N_),
633 make_embed_transform(make_tuple(Wo_), make_tuple(ConvStrideW_)),
634 make_pass_through_transform(NumGroupsToMerge),
635 make_pass_through_transform(C_)),
636 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
637 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
638
639 return transform_tensor_descriptor(
640 in_n_wo_c_desc,
641 make_tuple(make_merge_transform(make_tuple(N_, Wo_, NumGroupsToMerge)),
642 make_pass_through_transform(C_)),
643 make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}),
644 make_tuple(Sequence<0>{}, Sequence<1>{}));
645 }
646 }
647 else
648 {
649 if constexpr(NumGroupsToMerge == 1)
650 {
651 const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
652 make_tuple(N_, Wi_, C_),
653 make_tuple(NStrideTensorA_, WiStride_, CStrideTensorA_));
654
655 const auto in_n_wip_c_desc = transform_tensor_descriptor(
656 in_n_wi_c_desc,
657 make_tuple(make_pass_through_transform(N_),
658 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
659 make_pass_through_transform(C_)),
660 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
661 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
662
663 const auto in_n_x_wo_c_desc = transform_tensor_descriptor(
664 in_n_wip_c_desc,
665 make_tuple(make_pass_through_transform(N_),
666 make_embed_transform(make_tuple(X_, Wo_),
667 make_tuple(ConvDilationW_, ConvStrideW_)),
668 make_pass_through_transform(C_)),
669 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
670 make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}));
671
672 return transform_tensor_descriptor(
673 in_n_x_wo_c_desc,
674 make_tuple(make_merge_transform(make_tuple(N_, Wo_)),
675 make_merge_transform(make_tuple(X_, C_))),
676 make_tuple(Sequence<0, 2>{}, Sequence<1, 3>{}),
677 make_tuple(Sequence<0>{}, Sequence<1>{}));
678 }
679 else
680 {
681 const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
682 make_tuple(N_, Wi_, NumGroupsToMerge, C_),
683 make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_));
684
685 const auto in_n_wip_c_desc = transform_tensor_descriptor(
686 in_n_wi_c_desc,
687 make_tuple(make_pass_through_transform(N_),
688 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
689 make_pass_through_transform(NumGroupsToMerge),
690 make_pass_through_transform(C_)),
691 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
692 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
693
694 const auto in_n_x_wo_c_desc = transform_tensor_descriptor(
695 in_n_wip_c_desc,
696 make_tuple(make_pass_through_transform(N_),
697 make_embed_transform(make_tuple(X_, Wo_),
698 make_tuple(ConvDilationW_, ConvStrideW_)),
699 make_pass_through_transform(NumGroupsToMerge),
700 make_pass_through_transform(C_)),
701 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
702 make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}, Sequence<4>{}));
703
704 return transform_tensor_descriptor(
705 in_n_x_wo_c_desc,
706 make_tuple(make_merge_transform(make_tuple(N_, Wo_, NumGroupsToMerge)),
707 make_merge_transform(make_tuple(X_, C_))),
708 make_tuple(Sequence<0, 2, 3>{}, Sequence<1, 4>{}),
709 make_tuple(Sequence<0>{}, Sequence<1>{}));
710 }
711 }
712 }
713
714 template <typename ALayout,
715 typename ck::enable_if<NDimSpatial == 2 &&
716 (is_same_v<ALayout, tensor_layout::convolution::G_NHW_C> ||
717 is_same_v<ALayout, tensor_layout::convolution::NHWGC> ||
718 is_same_v<ALayout, tensor_layout::convolution::GNHWC>),
719 bool>::type = false>
720 __host__ __device__ auto MakeADescriptor_M_K() const
721
722 {
723 if constexpr(ConvForwardSpecialization ==
724 device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
725 {
726 if constexpr(NumGroupsToMerge == 1)
727 {
728 const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor(
729 make_tuple(N_, Ho_, Wo_, C_),
730 make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_));
731
732 return transform_tensor_descriptor(
733 in_gemmm_gemmk_desc,
734 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_)),
735 make_pass_through_transform(C_)),
736 make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}),
737 make_tuple(Sequence<0>{}, Sequence<1>{}));
738 }
739 else
740 {
741 const auto in_gemmm_groups_gemmk_desc = make_naive_tensor_descriptor(
742 make_tuple(N_, Ho_, Wo_, NumGroupsToMerge, C_),
743 make_tuple(
744 NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_));
745
746 return transform_tensor_descriptor(
747 in_gemmm_groups_gemmk_desc,
748 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_, NumGroupsToMerge)),
749 make_pass_through_transform(C_)),
750 make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}),
751 make_tuple(Sequence<0>{}, Sequence<1>{}));
752 }
753 }
754 else if constexpr(ConvForwardSpecialization ==
755 device::ConvolutionForwardSpecialization::Filter3x3)
756 {
757 if constexpr(NumGroupsToMerge == 1)
758 {
759 const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor(
760 make_tuple(N_, Hi_, Wi_), make_tuple(NStrideTensorA_, HiStride_, WiStride_));
761
762 const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
763 in_n_hi_wi_c_desc,
764 make_tuple(make_pass_through_transform(N_),
765 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
766 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_)),
767 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
768 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
769
770 const auto in_n_y_ho_x_wo_c_desc = transform_tensor_descriptor(
771 in_n_hip_wip_c_desc,
772 make_tuple(make_pass_through_transform(N_),
773 make_embed_transform(make_tuple(Number<3>{}, Ho_),
774 make_tuple(ConvDilationH_, ConvStrideH_)),
775 make_embed_transform(make_tuple(Number<3>{}, Wo_),
776 make_tuple(ConvDilationW_, ConvStrideW_))),
777 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
778 make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}));
779
780 return transform_tensor_descriptor(
781 in_n_y_ho_x_wo_c_desc,
782 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_)),
783 make_merge_transform(make_tuple(Number<3>{}, Number<3>{}))),
784 make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3>{}),
785 make_tuple(Sequence<0>{}, Sequence<1>{}));
786 }
787 else
788 {
789 const auto in_n_hi_wi_groups_c_desc = make_naive_tensor_descriptor(
790 make_tuple(N_, Hi_, Wi_, NumGroupsToMerge),
791 make_tuple(NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_));
792
793 const auto in_n_hip_wip_groups_c_desc = transform_tensor_descriptor(
794 in_n_hi_wi_groups_c_desc,
795 make_tuple(make_pass_through_transform(N_),
796 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
797 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
798 make_pass_through_transform(NumGroupsToMerge)),
799 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
800 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
801
802 const auto in_n_y_ho_x_wo_groups_c_desc = transform_tensor_descriptor(
803 in_n_hip_wip_groups_c_desc,
804 make_tuple(make_pass_through_transform(N_),
805 make_embed_transform(make_tuple(Number<3>{}, Ho_),
806 make_tuple(ConvDilationH_, ConvStrideH_)),
807 make_embed_transform(make_tuple(Number<3>{}, Wo_),
808 make_tuple(ConvDilationW_, ConvStrideW_)),
809 make_pass_through_transform(NumGroupsToMerge)),
810 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
811 make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
812
813 return transform_tensor_descriptor(
814 in_n_y_ho_x_wo_groups_c_desc,
815 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_, NumGroupsToMerge)),
816 make_merge_transform(make_tuple(Number<3>{}, Number<3>{}))),
817 make_tuple(Sequence<0, 2, 4, 5>{}, Sequence<1, 3>{}),
818 make_tuple(Sequence<0>{}, Sequence<1>{}));
819 }
820 }
821 else if constexpr(ConvForwardSpecialization ==
822 device::ConvolutionForwardSpecialization::Filter1x1Pad0)
823 {
824 if constexpr(NumGroupsToMerge == 1)
825 {
826 const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor(
827 make_tuple(N_, Hi_, Wi_, C_),
828 make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_));
829
830 const auto in_n_ho_wo_c_desc = transform_tensor_descriptor(
831 in_n_hi_wi_c_desc,
832 make_tuple(make_pass_through_transform(N_),
833 make_embed_transform(make_tuple(Ho_), make_tuple(ConvStrideH_)),
834 make_embed_transform(make_tuple(Wo_), make_tuple(ConvStrideW_)),
835 make_pass_through_transform(C_)),
836 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
837 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
838
839 return transform_tensor_descriptor(
840 in_n_ho_wo_c_desc,
841 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_)),
842 make_pass_through_transform(C_)),
843 make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}),
844 make_tuple(Sequence<0>{}, Sequence<1>{}));
845 }
846 else
847 {
848 const auto in_n_hi_wi_groups_c_desc = make_naive_tensor_descriptor(
849 make_tuple(N_, Hi_, Wi_, NumGroupsToMerge, C_),
850 make_tuple(
851 NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_));
852
853 const auto in_n_ho_wo_groups_c_desc = transform_tensor_descriptor(
854 in_n_hi_wi_groups_c_desc,
855 make_tuple(make_pass_through_transform(N_),
856 make_embed_transform(make_tuple(Ho_), make_tuple(ConvStrideH_)),
857 make_embed_transform(make_tuple(Wo_), make_tuple(ConvStrideW_)),
858 make_pass_through_transform(NumGroupsToMerge),
859 make_pass_through_transform(C_)),
860 make_tuple(
861 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
862 make_tuple(
863 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}));
864
865 return transform_tensor_descriptor(
866 in_n_ho_wo_groups_c_desc,
867 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_, NumGroupsToMerge)),
868 make_pass_through_transform(C_)),
869 make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}),
870 make_tuple(Sequence<0>{}, Sequence<1>{}));
871 }
872 }
873 else
874 {
875 if constexpr(NumGroupsToMerge == 1)
876 {
877 const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor(
878 make_tuple(N_, Hi_, Wi_, C_),
879 make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_));
880
881 const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
882 in_n_hi_wi_c_desc,
883 make_tuple(make_pass_through_transform(N_),
884 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
885 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
886 make_pass_through_transform(C_)),
887 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
888 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
889
890 const auto in_n_y_ho_x_wo_c_desc = transform_tensor_descriptor(
891 in_n_hip_wip_c_desc,
892 make_tuple(make_pass_through_transform(N_),
893 make_embed_transform(make_tuple(Y_, Ho_),
894 make_tuple(ConvDilationH_, ConvStrideH_)),
895 make_embed_transform(make_tuple(X_, Wo_),
896 make_tuple(ConvDilationW_, ConvStrideW_)),
897 make_pass_through_transform(C_)),
898 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
899 make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
900
901 return transform_tensor_descriptor(
902 in_n_y_ho_x_wo_c_desc,
903 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_)),
904 make_merge_transform(make_tuple(Y_, X_, C_))),
905 make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3, 5>{}),
906 make_tuple(Sequence<0>{}, Sequence<1>{}));
907 }
908 else
909 {
910
911 const auto in_n_hi_wi_groups_c_desc = make_naive_tensor_descriptor(
912 make_tuple(N_, Hi_, Wi_, NumGroupsToMerge, C_),
913 make_tuple(
914 NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_));
915
916 const auto in_n_hip_wip_groups_c_desc = transform_tensor_descriptor(
917 in_n_hi_wi_groups_c_desc,
918 make_tuple(make_pass_through_transform(N_),
919 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
920 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
921 make_pass_through_transform(NumGroupsToMerge),
922 make_pass_through_transform(C_)),
923 make_tuple(
924 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
925 make_tuple(
926 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}));
927
928 const auto in_n_y_ho_x_wo_groups_c_desc = transform_tensor_descriptor(
929 in_n_hip_wip_groups_c_desc,
930 make_tuple(make_pass_through_transform(N_),
931 make_embed_transform(make_tuple(Y_, Ho_),
932 make_tuple(ConvDilationH_, ConvStrideH_)),
933 make_embed_transform(make_tuple(X_, Wo_),
934 make_tuple(ConvDilationW_, ConvStrideW_)),
935 make_pass_through_transform(NumGroupsToMerge),
936 make_pass_through_transform(C_)),
937 make_tuple(
938 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
939 make_tuple(Sequence<0>{},
940 Sequence<1, 2>{},
941 Sequence<3, 4>{},
942 Sequence<5>{},
943 Sequence<6>{}));
944
945 return transform_tensor_descriptor(
946 in_n_y_ho_x_wo_groups_c_desc,
947 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_, NumGroupsToMerge)),
948 make_merge_transform(make_tuple(Y_, X_, C_))),
949 make_tuple(Sequence<0, 2, 4, 5>{}, Sequence<1, 3, 6>{}),
950 make_tuple(Sequence<0>{}, Sequence<1>{}));
951 }
952 }
953 }
954
955 template <typename ALayout,
956 typename ck::enable_if<
957 NDimSpatial == 3 && (is_same_v<ALayout, tensor_layout::convolution::G_NDHW_C> ||
958 is_same_v<ALayout, tensor_layout::convolution::NDHWGC> ||
959 is_same_v<ALayout, tensor_layout::convolution::GNDHWC>),
960 bool>::type = false>
961 __host__ __device__ auto MakeADescriptor_M_K() const
962
963 {
964 if constexpr(ConvForwardSpecialization ==
965 device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
966 {
967 if constexpr(NumGroupsToMerge == 1)
968 {
969 const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor(
970 make_tuple(N_, Do_, Ho_, Wo_, C_),
971 make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_));
972
973 return transform_tensor_descriptor(
974 in_gemmm_gemmk_desc,
975 make_tuple(make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_)),
976 make_pass_through_transform(C_)),
977 make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}),
978 make_tuple(Sequence<0>{}, Sequence<1>{}));
979 }
980 else
981 {
982 const auto in_gemmm_groups_gemmk_desc = make_naive_tensor_descriptor(
983 make_tuple(N_, Do_, Ho_, Wo_, NumGroupsToMerge, C_),
984 make_tuple(NStrideTensorA_,
985 DiStride_,
986 HiStride_,
987 WiStride_,
988 GStrideTensorA_,
989 CStrideTensorA_));
990
991 return transform_tensor_descriptor(
992 in_gemmm_groups_gemmk_desc,
993 make_tuple(
994 make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_, NumGroupsToMerge)),
995 make_pass_through_transform(C_)),
996 make_tuple(Sequence<0, 1, 2, 3, 4>{}, Sequence<5>{}),
997 make_tuple(Sequence<0>{}, Sequence<1>{}));
998 }
999 }
1000 else if constexpr(ConvForwardSpecialization ==
1001 device::ConvolutionForwardSpecialization::Filter3x3)
1002 {
1003 if constexpr(NumGroupsToMerge == 1)
1004 {
1005 const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
1006 make_tuple(N_, Di_, Hi_, Wi_),
1007 make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_));
1008
1009 const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
1010 in_n_di_hi_wi_c_desc,
1011 make_tuple(make_pass_through_transform(N_),
1012 make_pad_transform(Di_, InLeftPadD_, InRightPadD_),
1013 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
1014 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_)),
1015 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
1016 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
1017
1018 const auto in_n_z_do_y_ho_x_wo_c_desc = transform_tensor_descriptor(
1019 in_n_hip_wip_c_desc,
1020 make_tuple(make_pass_through_transform(N_),
1021 make_embed_transform(make_tuple(Number<3>{}, Do_),
1022 make_tuple(ConvDilationD_, ConvStrideD_)),
1023 make_embed_transform(make_tuple(Number<3>{}, Ho_),
1024 make_tuple(ConvDilationH_, ConvStrideH_)),
1025 make_embed_transform(make_tuple(Number<3>{}, Wo_),
1026 make_tuple(ConvDilationW_, ConvStrideW_))),
1027 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
1028 make_tuple(
1029 Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5, 6>{}));
1030
1031 return transform_tensor_descriptor(
1032 in_n_z_do_y_ho_x_wo_c_desc,
1033 make_tuple(
1034 make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_)),
1035 make_merge_transform(make_tuple(Number<3>{}, Number<3>{}, Number<3>{}))),
1036 make_tuple(Sequence<0, 2, 4, 6>{}, Sequence<1, 3, 5>{}),
1037 make_tuple(Sequence<0>{}, Sequence<1>{}));
1038 }
1039 else
1040 {
1041 const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
1042 make_tuple(N_, Di_, Hi_, Wi_, NumGroupsToMerge),
1043 make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, GStrideTensorA_));
1044
1045 const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
1046 in_n_di_hi_wi_c_desc,
1047 make_tuple(make_pass_through_transform(N_),
1048 make_pad_transform(Di_, InLeftPadD_, InRightPadD_),
1049 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
1050 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
1051 make_pass_through_transform(NumGroupsToMerge)),
1052 make_tuple(
1053 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
1054 make_tuple(
1055 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}));
1056
1057 const auto in_n_z_do_y_ho_x_wo_c_desc = transform_tensor_descriptor(
1058 in_n_hip_wip_c_desc,
1059 make_tuple(make_pass_through_transform(N_),
1060 make_embed_transform(make_tuple(Number<3>{}, Do_),
1061 make_tuple(ConvDilationD_, ConvStrideD_)),
1062 make_embed_transform(make_tuple(Number<3>{}, Ho_),
1063 make_tuple(ConvDilationH_, ConvStrideH_)),
1064 make_embed_transform(make_tuple(Number<3>{}, Wo_),
1065 make_tuple(ConvDilationW_, ConvStrideW_)),
1066 make_pass_through_transform(NumGroupsToMerge)),
1067 make_tuple(
1068 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
1069 make_tuple(Sequence<0>{},
1070 Sequence<1, 2>{},
1071 Sequence<3, 4>{},
1072 Sequence<5, 6>{},
1073 Sequence<7>{}));
1074
1075 return transform_tensor_descriptor(
1076 in_n_z_do_y_ho_x_wo_c_desc,
1077 make_tuple(
1078 make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_, NumGroupsToMerge)),
1079 make_merge_transform(make_tuple(Number<3>{}, Number<3>{}, Number<3>{}))),
1080 make_tuple(Sequence<0, 2, 4, 6, 7>{}, Sequence<1, 3, 5>{}),
1081 make_tuple(Sequence<0>{}, Sequence<1>{}));
1082 }
1083 }
1084 else if constexpr(ConvForwardSpecialization ==
1085 device::ConvolutionForwardSpecialization::Filter1x1Pad0)
1086 {
1087 if constexpr(NumGroupsToMerge == 1)
1088 {
1089 const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
1090 make_tuple(N_, Di_, Hi_, Wi_, C_),
1091 make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_));
1092
1093 const auto in_n_do_ho_wo_c_desc = transform_tensor_descriptor(
1094 in_n_di_hi_wi_c_desc,
1095 make_tuple(make_pass_through_transform(N_),
1096 make_embed_transform(make_tuple(Do_), make_tuple(ConvStrideD_)),
1097 make_embed_transform(make_tuple(Ho_), make_tuple(ConvStrideH_)),
1098 make_embed_transform(make_tuple(Wo_), make_tuple(ConvStrideW_)),
1099 make_pass_through_transform(C_)),
1100 make_tuple(
1101 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
1102 make_tuple(
1103 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}));
1104
1105 return transform_tensor_descriptor(
1106 in_n_do_ho_wo_c_desc,
1107 make_tuple(make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_)),
1108 make_pass_through_transform(C_)),
1109 make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}),
1110 make_tuple(Sequence<0>{}, Sequence<1>{}));
1111 }
1112 else
1113 {
1114 const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
1115 make_tuple(N_, Di_, Hi_, Wi_, NumGroupsToMerge, C_),
1116 make_tuple(NStrideTensorA_,
1117 DiStride_,
1118 HiStride_,
1119 WiStride_,
1120 GStrideTensorA_,
1121 CStrideTensorA_));
1122
1123 const auto in_n_do_ho_wo_c_desc = transform_tensor_descriptor(
1124 in_n_di_hi_wi_c_desc,
1125 make_tuple(make_pass_through_transform(N_),
1126 make_embed_transform(make_tuple(Do_), make_tuple(ConvStrideD_)),
1127 make_embed_transform(make_tuple(Ho_), make_tuple(ConvStrideH_)),
1128 make_embed_transform(make_tuple(Wo_), make_tuple(ConvStrideW_)),
1129 make_pass_through_transform(NumGroupsToMerge),
1130 make_pass_through_transform(C_)),
1131 make_tuple(Sequence<0>{},
1132 Sequence<1>{},
1133 Sequence<2>{},
1134 Sequence<3>{},
1135 Sequence<4>{},
1136 Sequence<5>{}),
1137 make_tuple(Sequence<0>{},
1138 Sequence<1>{},
1139 Sequence<2>{},
1140 Sequence<3>{},
1141 Sequence<4>{},
1142 Sequence<5>{}));
1143
1144 return transform_tensor_descriptor(
1145 in_n_do_ho_wo_c_desc,
1146 make_tuple(
1147 make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_, NumGroupsToMerge)),
1148 make_pass_through_transform(C_)),
1149 make_tuple(Sequence<0, 1, 2, 3, 4>{}, Sequence<5>{}),
1150 make_tuple(Sequence<0>{}, Sequence<1>{}));
1151 }
1152 }
1153 else
1154 {
1155 if constexpr(NumGroupsToMerge == 1)
1156 {
1157 const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
1158 make_tuple(N_, Di_, Hi_, Wi_, C_),
1159 make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_));
1160
1161 const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
1162 in_n_di_hi_wi_c_desc,
1163 make_tuple(make_pass_through_transform(N_),
1164 make_pad_transform(Di_, InLeftPadD_, InRightPadD_),
1165 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
1166 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
1167 make_pass_through_transform(C_)),
1168 make_tuple(
1169 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
1170 make_tuple(
1171 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}));
1172
1173 const auto in_n_z_do_y_ho_x_wo_c_desc = transform_tensor_descriptor(
1174 in_n_hip_wip_c_desc,
1175 make_tuple(make_pass_through_transform(N_),
1176 make_embed_transform(make_tuple(Z_, Do_),
1177 make_tuple(ConvDilationD_, ConvStrideD_)),
1178 make_embed_transform(make_tuple(Y_, Ho_),
1179 make_tuple(ConvDilationH_, ConvStrideH_)),
1180 make_embed_transform(make_tuple(X_, Wo_),
1181 make_tuple(ConvDilationW_, ConvStrideW_)),
1182 make_pass_through_transform(C_)),
1183 make_tuple(
1184 Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
1185 make_tuple(Sequence<0>{},
1186 Sequence<1, 2>{},
1187 Sequence<3, 4>{},
1188 Sequence<5, 6>{},
1189 Sequence<7>{}));
1190
1191 return transform_tensor_descriptor(
1192 in_n_z_do_y_ho_x_wo_c_desc,
1193 make_tuple(make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_)),
1194 make_merge_transform(make_tuple(Z_, Y_, X_, C_))),
1195 make_tuple(Sequence<0, 2, 4, 6>{}, Sequence<1, 3, 5, 7>{}),
1196 make_tuple(Sequence<0>{}, Sequence<1>{}));
1197 }
1198 else
1199 {
1200 const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
1201 make_tuple(N_, Di_, Hi_, Wi_, NumGroupsToMerge, C_),
1202 make_tuple(NStrideTensorA_,
1203 DiStride_,
1204 HiStride_,
1205 WiStride_,
1206 GStrideTensorA_,
1207 CStrideTensorA_));
1208
1209 const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
1210 in_n_di_hi_wi_c_desc,
1211 make_tuple(make_pass_through_transform(N_),
1212 make_pad_transform(Di_, InLeftPadD_, InRightPadD_),
1213 make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
1214 make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
1215 make_pass_through_transform(NumGroupsToMerge),
1216 make_pass_through_transform(C_)),
1217 make_tuple(Sequence<0>{},
1218 Sequence<1>{},
1219 Sequence<2>{},
1220 Sequence<3>{},
1221 Sequence<4>{},
1222 Sequence<5>{}),
1223 make_tuple(Sequence<0>{},
1224 Sequence<1>{},
1225 Sequence<2>{},
1226 Sequence<3>{},
1227 Sequence<4>{},
1228 Sequence<5>{}));
1229
1230 const auto in_n_z_do_y_ho_x_wo_c_desc = transform_tensor_descriptor(
1231 in_n_hip_wip_c_desc,
1232 make_tuple(make_pass_through_transform(N_),
1233 make_embed_transform(make_tuple(Z_, Do_),
1234 make_tuple(ConvDilationD_, ConvStrideD_)),
1235 make_embed_transform(make_tuple(Y_, Ho_),
1236 make_tuple(ConvDilationH_, ConvStrideH_)),
1237 make_embed_transform(make_tuple(X_, Wo_),
1238 make_tuple(ConvDilationW_, ConvStrideW_)),
1239 make_pass_through_transform(NumGroupsToMerge),
1240 make_pass_through_transform(C_)),
1241 make_tuple(Sequence<0>{},
1242 Sequence<1>{},
1243 Sequence<2>{},
1244 Sequence<3>{},
1245 Sequence<4>{},
1246 Sequence<5>{}),
1247 make_tuple(Sequence<0>{},
1248 Sequence<1, 2>{},
1249 Sequence<3, 4>{},
1250 Sequence<5, 6>{},
1251 Sequence<7>{},
1252 Sequence<8>{}));
1253
1254 return transform_tensor_descriptor(
1255 in_n_z_do_y_ho_x_wo_c_desc,
1256 make_tuple(
1257 make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_, NumGroupsToMerge)),
1258 make_merge_transform(make_tuple(Z_, Y_, X_, C_))),
1259 make_tuple(Sequence<0, 2, 4, 6, 7>{}, Sequence<1, 3, 5, 8>{}),
1260 make_tuple(Sequence<0>{}, Sequence<1>{}));
1261 }
1262 }
1263 }
1264
1265 template <typename ALayout,
1266 typename ck::enable_if<NDimSpatial == 1 &&
1267 is_same_v<ALayout, tensor_layout::convolution::NGCW>,
1268 bool>::type = false>
1269 __host__ __device__ auto MakeADescriptor_M_K() const
1270 {
1271 static_assert(NumGroupsToMerge == 1);
1272 static_assert(ConvForwardSpecialization ==
1273 device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0);
1274
1275 const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor(
1276 make_tuple(N_, Wo_, C_), make_tuple(NStrideTensorA_, I1, CStrideTensorA_));
1277
1278 return transform_tensor_descriptor(
1279 in_gemmm_gemmk_desc,
1280 make_tuple(make_merge_transform(make_tuple(N_, Wo_)), make_pass_through_transform(C_)),
1281 make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
1282 make_tuple(Sequence<0>{}, Sequence<1>{}));
1283 }
1284
1285 template <typename ALayout,
1286 typename ck::enable_if<NDimSpatial == 2 &&
1287 is_same_v<ALayout, tensor_layout::convolution::NGCHW>,
1288 bool>::type = false>
1289 __host__ __device__ auto MakeADescriptor_M_K() const
1290 {
1291 static_assert(NumGroupsToMerge == 1);
1292 static_assert(ConvForwardSpecialization ==
1293 device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0);
1294
1295 const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor(
1296 make_tuple(N_, Ho_ * Wo_, C_), make_tuple(NStrideTensorA_, I1, CStrideTensorA_));
1297
1298 return transform_tensor_descriptor(
1299 in_gemmm_gemmk_desc,
1300 make_tuple(make_merge_transform(make_tuple(N_, Ho_ * Wo_)),
1301 make_pass_through_transform(C_)),
1302 make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
1303 make_tuple(Sequence<0>{}, Sequence<1>{}));
1304 }
1305
1306 template <typename ALayout,
1307 typename ck::enable_if<NDimSpatial == 3 &&
1308 is_same_v<ALayout, tensor_layout::convolution::NGCDHW>,
1309 bool>::type = false>
1310 __host__ __device__ auto MakeADescriptor_M_K() const
1311 {
1312 static_assert(NumGroupsToMerge == 1);
1313 static_assert(ConvForwardSpecialization ==
1314 device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0);
1315
1316 const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor(
1317 make_tuple(N_, Do_ * Ho_ * Wo_, C_), make_tuple(NStrideTensorA_, I1, CStrideTensorA_));
1318
1319 return transform_tensor_descriptor(
1320 in_gemmm_gemmk_desc,
1321 make_tuple(make_merge_transform(make_tuple(N_, Do_ * Ho_ * Wo_)),
1322 make_pass_through_transform(C_)),
1323 make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
1324 make_tuple(Sequence<0>{}, Sequence<1>{}));
1325 }
1326
1327 template <typename BLayout,
1328 typename ck::enable_if<is_same_v<BLayout, tensor_layout::convolution::GKCX> ||
1329 is_same_v<BLayout, tensor_layout::convolution::GKCYX> ||
1330 is_same_v<BLayout, tensor_layout::convolution::GKCZYX>,
1331 bool>::type = false>
1332 __host__ __device__ auto MakeBDescriptor_N_K() const
1333 {
1334 static_assert(ConvForwardSpecialization ==
1335 device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0 ||
1336 ConvForwardSpecialization ==
1337 device::ConvolutionForwardSpecialization::Filter1x1Pad0);
1338 static_assert(NumGroupsToMerge == 1);
1339 return make_naive_tensor_descriptor_packed(make_tuple(K_, C_));
1340 }
1341
1342 template <typename BLayout,
1343 typename ck::enable_if<is_same_v<BLayout, tensor_layout::convolution::GKXC> ||
1344 is_same_v<BLayout, tensor_layout::convolution::GKYXC> ||
1345 is_same_v<BLayout, tensor_layout::convolution::GKZYXC>,
1346 bool>::type = false>
1347 __host__ __device__ auto MakeBDescriptor_N_K() const
1348 {
1349 if constexpr(ConvForwardSpecialization ==
1350 device::ConvolutionForwardSpecialization::Filter3x3)
1351 {
1352 using FilterSizeNumType =
1353 ck::conditional_t<NDimSpatial == 1,
1354 Number<3>,
1355 ck::conditional_t<NDimSpatial == 2, Number<9>, Number<27>>>;
1356
1357 if constexpr(NumGroupsToMerge == 1)
1358 {
1359 return make_naive_tensor_descriptor_packed(make_tuple(K_, FilterSizeNumType{}));
1360 }
1361 else
1362 {
1363
1364 const auto wei_gemmn_groups_gemmk_desc = make_naive_tensor_descriptor(
1365 make_tuple(K_, NumGroupsToMerge, FilterSizeNumType{}),
1366 make_tuple(KStrideTensorB_, GStrideTensorB_, CStrideTensorB_));
1367 return transform_tensor_descriptor(
1368 wei_gemmn_groups_gemmk_desc,
1369 make_tuple(make_merge_transform(make_tuple(K_, NumGroupsToMerge)),
1370 make_pass_through_transform(FilterSizeNumType{})),
1371 make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
1372 make_tuple(Sequence<0>{}, Sequence<1>{}));
1373 }
1374 }
1375 else
1376 {
1377 if constexpr(NumGroupsToMerge == 1)
1378 {
1379 return make_naive_tensor_descriptor_packed(make_tuple(K_, ZYX_ * C_));
1380 }
1381 else
1382 {
1383 const auto wei_gemmn_groups_gemmk_desc = make_naive_tensor_descriptor(
1384 make_tuple(K_, NumGroupsToMerge, ZYX_ * C_),
1385 make_tuple(KStrideTensorB_, GStrideTensorB_, CStrideTensorB_));
1386 return transform_tensor_descriptor(
1387 wei_gemmn_groups_gemmk_desc,
1388 make_tuple(make_merge_transform(make_tuple(K_, NumGroupsToMerge)),
1389 make_pass_through_transform(ZYX_ * C_)),
1390 make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
1391 make_tuple(Sequence<0>{}, Sequence<1>{}));
1392 }
1393 }
1394 }
1395
1396 template <
1397 typename BLayout,
1398 typename ck::enable_if<is_same_v<BLayout, tensor_layout::convolution::G_K_X_C> ||
1399 is_same_v<BLayout, tensor_layout::convolution::G_K_YX_C> ||
1400 is_same_v<BLayout, tensor_layout::convolution::G_K_ZYX_C> ||
1401 is_same_v<BLayout, tensor_layout::convolution::KXGC> ||
1402 is_same_v<BLayout, tensor_layout::convolution::KYXGC> ||
1403 is_same_v<BLayout, tensor_layout::convolution::KZYXGC>,
1404 bool>::type = false>
1405 __host__ __device__ auto MakeBDescriptor_N_K() const
1406 {
1407 const auto wei_k_yx_c_desc = make_naive_tensor_descriptor(
1408 make_tuple(K_, ZYX_, C_), make_tuple(KStrideTensorB_, XStride_, CStrideTensorB_));
1409
1410 const auto wei_gemmn_gemmk_desc = transform_tensor_descriptor(
1411 wei_k_yx_c_desc,
1412 make_tuple(make_pass_through_transform(K_), make_merge_transform(make_tuple(ZYX_, C_))),
1413 make_tuple(Sequence<0>{}, Sequence<1, 2>{}),
1414 make_tuple(Sequence<0>{}, Sequence<1>{}));
1415
1416 return wei_gemmn_gemmk_desc;
1417 }
1418
1419 template <
1420 typename CLayout,
1421 index_t NDimSp = NDimSpatial,
1422
1423 typename ck::enable_if<NDimSp == 1 && (is_same_v<CLayout, tensor_layout::convolution::G_K>),
1424 bool>::type = false>
1425 __host__ __device__ auto MakeCDescriptor_M_N() const
1426 {
1427 if constexpr(CTranspose)
1428 {
1429 return make_naive_tensor_descriptor(make_tuple(K_, N_ * Wo_),
1430 make_tuple(KStrideTensorC_, I0));
1431 }
1432 else
1433 {
1434 return make_naive_tensor_descriptor(make_tuple(N_ * Wo_, K_),
1435 make_tuple(I0, KStrideTensorC_));
1436 }
1437 }
1438
1439 template <
1440 typename CLayout,
1441 index_t NDimSp = NDimSpatial,
1442
1443 typename ck::enable_if<NDimSp == 2 && (is_same_v<CLayout, tensor_layout::convolution::G_K>),
1444 bool>::type = false>
1445 __host__ __device__ auto MakeCDescriptor_M_N() const
1446 {
1447 if constexpr(CTranspose)
1448 {
1449 return make_naive_tensor_descriptor(make_tuple(K_, N_ * Ho_ * Wo_),
1450 make_tuple(KStrideTensorC_, I0));
1451 }
1452 else
1453 {
1454 return make_naive_tensor_descriptor(make_tuple(N_ * Ho_ * Wo_, K_),
1455 make_tuple(I0, KStrideTensorC_));
1456 }
1457 }
1458
1459 template <
1460 typename CLayout,
1461 index_t NDimSp = NDimSpatial,
1462
1463 typename ck::enable_if<NDimSp == 3 && (is_same_v<CLayout, tensor_layout::convolution::G_K>),
1464 bool>::type = false>
1465 __host__ __device__ auto MakeCDescriptor_M_N() const
1466 {
1467 if constexpr(CTranspose)
1468 {
1469 return make_naive_tensor_descriptor(make_tuple(K_, N_ * Do_ * Ho_ * Wo_),
1470 make_tuple(KStrideTensorC_, I0));
1471 }
1472 else
1473 {
1474 return make_naive_tensor_descriptor(make_tuple(N_ * Do_ * Ho_ * Wo_, K_),
1475 make_tuple(I0, KStrideTensorC_));
1476 }
1477 }
1478
1479 template <typename CLayout,
1480 index_t NDimSp = NDimSpatial,
1481
1482 typename ck::enable_if<NDimSp == 1 &&
1483 (is_same_v<CLayout, tensor_layout::convolution::G_NW_K> ||
1484 is_same_v<CLayout, tensor_layout::convolution::NWGK> ||
1485 is_same_v<CLayout, tensor_layout::convolution::GNWK>),
1486 bool>::type = false>
1487 __host__ __device__ auto MakeCDescriptor_M_N() const
1488 {
1489 static_assert(CTranspose == false);
1490 const IndexType NDoHoWo = N_ * Wo_;
1491 if constexpr(NumGroupsToMerge == 1)
1492 {
1493 return make_naive_tensor_descriptor(make_tuple(NDoHoWo, K_),
1494 make_tuple(WoStride_, KStrideTensorC_));
1495 }
1496 else
1497 {
1498 const auto nhwo_groups_k_1_desc = make_naive_tensor_descriptor(
1499 make_tuple(N_, Wo_, NumGroupsToMerge, K_, 1),
1500 make_tuple(
1501 NStrideTensorC_, WoStride_, GStrideTensorC_, KStrideTensorC_, GStrideTensorC_));
1502 // Padd 1 to NumGroupsToMerge
1503 const auto padded_desc = transform_tensor_descriptor(
1504 nhwo_groups_k_1_desc,
1505 make_tuple(make_merge_transform(make_tuple(N_, Wo_)),
1506 make_pass_through_transform(NumGroupsToMerge),
1507 make_pass_through_transform(K_),
1508 make_pad_transform(1, 0, NumGroupsToMerge - 1)),
1509 make_tuple(Sequence<0, 1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
1510 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
1511 // We need only matrices from diagonal. X_or returns 0 for the same
1512 // values. So if matrices is not on diagonal then it will be stored in padding.
1513 // To avoid use of modulo after xor we assume that NumBatch to merge is power of 2.
1514 static_assert(NumGroupsToMerge == 1 || NumGroupsToMerge == 2 || NumGroupsToMerge == 4 ||
1515 NumGroupsToMerge == 8 || NumGroupsToMerge == 16 ||
1516 NumGroupsToMerge == 32 || NumGroupsToMerge == 64);
1517 const auto unmerged_padded_desc = transform_tensor_descriptor(
1518 padded_desc,
1519 make_tuple(make_pass_through_transform(NDoHoWo),
1520 make_xor_transform(make_tuple(NumGroupsToMerge, NumGroupsToMerge)),
1521 make_pass_through_transform(K_)),
1522 make_tuple(Sequence<0>{}, Sequence<1, 3>{}, Sequence<2>{}),
1523 make_tuple(Sequence<0>{}, Sequence<1, 3>{}, Sequence<2>{}));
1524 // Merge To M, N
1525 return transform_tensor_descriptor(
1526 unmerged_padded_desc,
1527 make_tuple(make_merge_transform(make_tuple(NDoHoWo, NumGroupsToMerge)),
1528 make_merge_transform(make_tuple(K_, NumGroupsToMerge))),
1529 make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{}),
1530 make_tuple(Sequence<0>{}, Sequence<1>{}));
1531 }
1532 }
1533
1534 template <typename CLayout,
1535 index_t NDimSp = NDimSpatial,
1536
1537 typename ck::enable_if<NDimSp == 2 &&
1538 (is_same_v<CLayout, tensor_layout::convolution::G_NHW_K> ||
1539 is_same_v<CLayout, tensor_layout::convolution::NHWGK> ||
1540 is_same_v<CLayout, tensor_layout::convolution::GNHWK>),
1541 bool>::type = false>
1542 __host__ __device__ auto MakeCDescriptor_M_N() const
1543 {
1544 static_assert(CTranspose == false);
1545 const IndexType NDoHoWo = N_ * Ho_ * Wo_;
1546 if constexpr(NumGroupsToMerge == 1)
1547 {
1548 return make_naive_tensor_descriptor(make_tuple(NDoHoWo, K_),
1549 make_tuple(WoStride_, KStrideTensorC_));
1550 }
1551 else
1552 {
1553 const auto nhwo_groups_k_1_desc =
1554 make_naive_tensor_descriptor(make_tuple(N_, Ho_, Wo_, NumGroupsToMerge, K_, 1),
1555 make_tuple(NStrideTensorC_,
1556 HoStride_,
1557 WoStride_,
1558 GStrideTensorC_,
1559 KStrideTensorC_,
1560 GStrideTensorC_));
1561 // Padd 1 to NumGroupsToMerge
1562 const auto padded_desc = transform_tensor_descriptor(
1563 nhwo_groups_k_1_desc,
1564 make_tuple(make_merge_transform(make_tuple(N_, Ho_, Wo_)),
1565 make_pass_through_transform(NumGroupsToMerge),
1566 make_pass_through_transform(K_),
1567 make_pad_transform(1, 0, NumGroupsToMerge - 1)),
1568 make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}, Sequence<4>{}, Sequence<5>{}),
1569 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
1570 // We need only matrices from diagonal. X_or returns 0 for the same
1571 // values. So if matrices is not on diagonal then it will be stored in padding.
1572 // To avoid use of modulo after xor we assume that NumBatch to merge is power of 2.
1573 static_assert(NumGroupsToMerge == 1 || NumGroupsToMerge == 2 || NumGroupsToMerge == 4 ||
1574 NumGroupsToMerge == 8 || NumGroupsToMerge == 16 ||
1575 NumGroupsToMerge == 32 || NumGroupsToMerge == 64);
1576 const auto unmerged_padded_desc = transform_tensor_descriptor(
1577 padded_desc,
1578 make_tuple(make_pass_through_transform(NDoHoWo),
1579 make_xor_transform(make_tuple(NumGroupsToMerge, NumGroupsToMerge)),
1580 make_pass_through_transform(K_)),
1581 make_tuple(Sequence<0>{}, Sequence<1, 3>{}, Sequence<2>{}),
1582 make_tuple(Sequence<0>{}, Sequence<1, 3>{}, Sequence<2>{}));
1583 // Merge To M, N
1584 return transform_tensor_descriptor(
1585 unmerged_padded_desc,
1586 make_tuple(make_merge_transform(make_tuple(NDoHoWo, NumGroupsToMerge)),
1587 make_merge_transform(make_tuple(K_, NumGroupsToMerge))),
1588 make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{}),
1589 make_tuple(Sequence<0>{}, Sequence<1>{}));
1590 }
1591 }
1592
1593 template <typename CLayout,
1594 index_t NDimSp = NDimSpatial,
1595 typename ck::enable_if<
1596 NDimSp == 3 && (is_same_v<CLayout, tensor_layout::convolution::G_NDHW_K> ||
1597 is_same_v<CLayout, tensor_layout::convolution::NDHWGK> ||
1598 is_same_v<CLayout, tensor_layout::convolution::GNDHWK>),
1599 bool>::type = false>
1600 __host__ __device__ auto MakeCDescriptor_M_N() const
1601 {
1602 static_assert(CTranspose == false);
1603 const IndexType NDoHoWo = N_ * Do_ * Ho_ * Wo_;
1604 if constexpr(NumGroupsToMerge == 1)
1605 {
1606 return make_naive_tensor_descriptor(make_tuple(NDoHoWo, K_),
1607 make_tuple(WoStride_, KStrideTensorC_));
1608 }
1609 else
1610 {
1611 const auto nhwo_groups_k_1_desc =
1612 make_naive_tensor_descriptor(make_tuple(N_, Do_, Ho_, Wo_, NumGroupsToMerge, K_, 1),
1613 make_tuple(NStrideTensorC_,
1614 DoStride_,
1615 HoStride_,
1616 WoStride_,
1617 GStrideTensorC_,
1618 KStrideTensorC_,
1619 GStrideTensorC_));
1620 // Padd 1 to NumGroupsToMerge
1621 const auto padded_desc = transform_tensor_descriptor(
1622 nhwo_groups_k_1_desc,
1623 make_tuple(make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_)),
1624 make_pass_through_transform(NumGroupsToMerge),
1625 make_pass_through_transform(K_),
1626 make_pad_transform(1, 0, NumGroupsToMerge - 1)),
1627 make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}, Sequence<5>{}, Sequence<6>{}),
1628 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
1629 // We need only matrices from diagonal. X_or returns 0 for the same
1630 // values. So if matrices is not on diagonal then it will be stored in padding.
1631 // To avoid use of modulo after xor we assume that NumBatch to merge is power of 2.
1632 static_assert(NumGroupsToMerge == 1 || NumGroupsToMerge == 2 || NumGroupsToMerge == 4 ||
1633 NumGroupsToMerge == 8 || NumGroupsToMerge == 16 ||
1634 NumGroupsToMerge == 32 || NumGroupsToMerge == 64);
1635 const auto unmerged_padded_desc = transform_tensor_descriptor(
1636 padded_desc,
1637 make_tuple(make_pass_through_transform(NDoHoWo),
1638 make_xor_transform(make_tuple(NumGroupsToMerge, NumGroupsToMerge)),
1639 make_pass_through_transform(K_)),
1640 make_tuple(Sequence<0>{}, Sequence<1, 3>{}, Sequence<2>{}),
1641 make_tuple(Sequence<0>{}, Sequence<1, 3>{}, Sequence<2>{}));
1642 // Merge To M, N
1643 return transform_tensor_descriptor(
1644 unmerged_padded_desc,
1645 make_tuple(make_merge_transform(make_tuple(NDoHoWo, NumGroupsToMerge)),
1646 make_merge_transform(make_tuple(K_, NumGroupsToMerge))),
1647 make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{}),
1648 make_tuple(Sequence<0>{}, Sequence<1>{}));
1649 }
1650 }
1651
1652 template <typename CLayout,
1653 index_t NDimSp = NDimSpatial,
1654
1655 typename ck::enable_if<NDimSp == 1 &&
1656 (is_same_v<CLayout, tensor_layout::convolution::GNKW> ||
1657 is_same_v<CLayout, tensor_layout::convolution::NGKW>),
1658 bool>::type = false>
1659 __host__ __device__ auto MakeCDescriptor_M_N() const
1660 {
1661 static_assert(NumGroupsToMerge == 1);
1662 auto n_k_wo_desc = make_naive_tensor_descriptor(
1663 make_tuple(N_, K_, Wo_), make_tuple(NStrideTensorC_, KStrideTensorC_, I1));
1664 if constexpr(CTranspose)
1665 {
1666 return transform_tensor_descriptor(
1667 n_k_wo_desc,
1668 make_tuple(make_pass_through_transform(K_),
1669 make_merge_transform(make_tuple(N_, Wo_))),
1670 make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
1671 make_tuple(Sequence<0>{}, Sequence<1>{}));
1672 }
1673 else
1674 {
1675 return transform_tensor_descriptor(n_k_wo_desc,
1676 make_tuple(make_merge_transform(make_tuple(N_, Wo_)),
1677 make_pass_through_transform(K_)),
1678 make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
1679 make_tuple(Sequence<0>{}, Sequence<1>{}));
1680 }
1681 }
1682
1683 template <typename CLayout,
1684 index_t NDimSp = NDimSpatial,
1685
1686 typename ck::enable_if<NDimSp == 2 &&
1687 (is_same_v<CLayout, tensor_layout::convolution::GNKHW> ||
1688 is_same_v<CLayout, tensor_layout::convolution::NGKHW>),
1689 bool>::type = false>
1690 __host__ __device__ auto MakeCDescriptor_M_N() const
1691 {
1692 static_assert(NumGroupsToMerge == 1);
1693 auto n_k_howo_desc = make_naive_tensor_descriptor(
1694 make_tuple(N_, K_, Ho_ * Wo_), make_tuple(NStrideTensorC_, KStrideTensorC_, I1));
1695 if constexpr(CTranspose)
1696 {
1697 return transform_tensor_descriptor(
1698 n_k_howo_desc,
1699 make_tuple(make_pass_through_transform(K_),
1700 make_merge_transform(make_tuple(N_, Ho_ * Wo_))),
1701 make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
1702 make_tuple(Sequence<0>{}, Sequence<1>{}));
1703 }
1704 else
1705 {
1706 return transform_tensor_descriptor(
1707 n_k_howo_desc,
1708 make_tuple(make_merge_transform(make_tuple(N_, Ho_ * Wo_)),
1709 make_pass_through_transform(K_)),
1710 make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
1711 make_tuple(Sequence<0>{}, Sequence<1>{}));
1712 }
1713 }
1714
1715 template <typename CLayout,
1716 index_t NDimSp = NDimSpatial,
1717
1718 typename ck::enable_if<NDimSp == 3 &&
1719 (is_same_v<CLayout, tensor_layout::convolution::GNKDHW> ||
1720 is_same_v<CLayout, tensor_layout::convolution::NGKDHW>),
1721 bool>::type = false>
1722 __host__ __device__ auto MakeCDescriptor_M_N() const
1723 {
1724 static_assert(NumGroupsToMerge == 1);
1725 auto n_k_dohowo_desc = make_naive_tensor_descriptor(
1726 make_tuple(N_, K_, Do_ * Ho_ * Wo_), make_tuple(NStrideTensorC_, KStrideTensorC_, I1));
1727
1728 if constexpr(CTranspose)
1729 {
1730 return transform_tensor_descriptor(
1731 n_k_dohowo_desc,
1732 make_tuple(make_pass_through_transform(K_),
1733 make_merge_transform(make_tuple(N_, Do_ * Ho_ * Wo_))),
1734 make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
1735 make_tuple(Sequence<0>{}, Sequence<1>{}));
1736 }
1737 else
1738 {
1739 return transform_tensor_descriptor(
1740 n_k_dohowo_desc,
1741 make_tuple(make_merge_transform(make_tuple(N_, Do_ * Ho_ * Wo_)),
1742 make_pass_through_transform(K_)),
1743 make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
1744 make_tuple(Sequence<0>{}, Sequence<1>{}));
1745 }
1746 }
1747 IndexType N_;
1748 IndexType Di_, Hi_, Wi_;
1749 IndexType Do_, Ho_, Wo_;
1750 IndexType Z_, Y_, X_;
1751 IndexType K_, C_;
1752 IndexType DiStride_, HiStride_, WiStride_;
1753 IndexType DoStride_, HoStride_, WoStride_;
1754 IndexType XStride_;
1755 IndexType CStrideTensorA_, CStrideTensorB_, KStrideTensorB_, KStrideTensorC_;
1756 IndexType NStrideTensorA_, NStrideTensorC_;
1757 IndexType GStrideTensorA_, GStrideTensorB_, GStrideTensorC_;
1758 IndexType ConvStrideD_, ConvStrideH_, ConvStrideW_;
1759 IndexType ConvDilationD_, ConvDilationH_, ConvDilationW_;
1760 IndexType InLeftPadD_, InLeftPadH_, InLeftPadW_;
1761 IndexType InRightPadD_, InRightPadH_, InRightPadW_;
1762 IndexType ZYX_;
1763};
1764
1765// wrapper class to call member functions on TransformConvToGemm struct at runtime
1766// TODO: figure out aq way to properly pass in layout as an argument
1767struct TransformConv
1768{
1769 TransformConv() {}
1770
1771 template <index_t NDimSpatial,
1772 device::ConvolutionForwardSpecialization ConvForwardSpecialization>
1773 auto
1774 transform_func(TransformConvFwdToGemm<NDimSpatial, ConvForwardSpecialization> conv_fwd_to_gemm)
1775 {
1776 if(NDimSpatial == 2)
1777 {
1778 return conv_fwd_to_gemm
1779 .template MakeCDescriptor_M_N<ck::tensor_layout::convolution::NHWGK, 2>();
1780 }
1781 else if(NDimSpatial == 3)
1782 {
1783 return conv_fwd_to_gemm
1784 .template MakeCDescriptor_M_N<tensor_layout::convolution::NDHWGK, 3>();
1785 }
1786 else if(NDimSpatial == 1)
1787 {
1788 return conv_fwd_to_gemm
1789 .template MakeCDescriptor_M_N<tensor_layout::convolution::NWGK, 1>();
1790 }
1791 }
1792};
1793
1794} // namespace tensor_operation
1795} // namespace ck
common_header.hpp
convolution_forward_specialization.hpp
ck::math::max
__host__ __device__ constexpr T max(T x)
Definition utility/math.hpp:84
ck::math::integer_divide_ceil
__host__ __device__ constexpr auto integer_divide_ceil(X x, Y y)
Definition utility/math.hpp:72
ck::math::min
__host__ __device__ constexpr T min(T x)
Definition utility/math.hpp:116
ck::tensor_operation::device::ConvolutionForwardSpecialization
ConvolutionForwardSpecialization
Definition convolution_forward_specialization.hpp:15
ck::tensor_operation::device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0
@ Filter1x1Stride1Pad0
Definition convolution_forward_specialization.hpp:18
ck::tensor_operation::device::ConvolutionForwardSpecialization::Filter3x3
@ Filter3x3
Definition convolution_forward_specialization.hpp:20
ck::tensor_operation::device::ConvolutionForwardSpecialization::Filter1x1Pad0
@ Filter1x1Pad0
Definition convolution_forward_specialization.hpp:17
ck::tensor_operation
Definition convolution_backward_data_specialization.hpp:7
ck
Definition ck.hpp:268
ck::make_pass_through_transform
__host__ __device__ constexpr auto make_pass_through_transform(const LowLength &low_length)
Definition multi_index_transform_helper.hpp:12
ck::index_t
int32_t index_t
Definition ck.hpp:299
ck::conditional_t
typename conditional< predicate, X, Y >::type conditional_t
Definition utility/functional.hpp:115
ck::make_naive_tensor_descriptor
__host__ __device__ constexpr auto make_naive_tensor_descriptor(const Tuple< Lengths... > &lengths, const Tuple< Strides... > &strides)
Definition tensor_descriptor_helper.hpp:49
ck::Number
integral_constant< index_t, N > Number
Definition number.hpp:12
ck::make_pad_transform
__host__ __device__ constexpr auto make_pad_transform(const LowLength &low_length, const LeftPad &left_pad, const RightPad &right_pad, integral_constant< bool, SkipIsValidCheck >=integral_constant< bool, false >{})
Definition multi_index_transform_helper.hpp:19
ck::make_embed_transform
__host__ __device__ constexpr auto make_embed_transform(const UpLengths &up_lengths, const Coefficients &coefficients)
Definition multi_index_transform_helper.hpp:48
ck::make_merge_transform
__host__ __device__ constexpr auto make_merge_transform(const LowLengths &low_lengths)
Definition multi_index_transform_helper.hpp:55
ck::enable_if
std::enable_if< B, T > enable_if
Definition enable_if.hpp:24
ck::is_same_v
constexpr bool is_same_v
Definition type.hpp:283
ck::generate_tuple
__host__ __device__ constexpr auto generate_tuple(F &&f, Number< N >)
Definition tuple_helper.hpp:21
ck::make_naive_tensor_descriptor_packed
__host__ __device__ constexpr auto make_naive_tensor_descriptor_packed(const Tuple< Lengths... > &lengths)
Definition tensor_descriptor_helper.hpp:101
ck::make_tuple
__host__ __device__ constexpr auto make_tuple(Xs &&... xs)
Definition utility/tuple.hpp:211
ck::long_index_t
int64_t long_index_t
Definition ck.hpp:300
ck::transform_tensor_descriptor
__host__ __device__ constexpr auto transform_tensor_descriptor(const OldTensorDescriptor &old_tensor_desc, const NewTransforms &new_transforms, NewLowerDimensionOldVisibleIdss, NewUpperDimensionNewVisibleIdss)
Definition tensor_description/tensor_descriptor.hpp:319
ck::make_xor_transform
__host__ __device__ constexpr auto make_xor_transform(const LowLengths &low_lengths)
Definition multi_index_transform_helper.hpp:191
ck::Sequence
Definition utility/sequence.hpp:43
ck::tensor_operation::TransformConvFwdToGemm
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:25
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Di_
index_t Di_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1748
ck::tensor_operation::TransformConvFwdToGemm::MakeBDescriptor_N_K
__host__ __device__ auto MakeBDescriptor_N_K() const
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1332
ck::tensor_operation::TransformConvFwdToGemm::TransformConvFwdToGemm
__host__ __device__ constexpr TransformConvFwdToGemm()
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:101
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::InRightPadH_
index_t InRightPadH_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1761
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::KStrideTensorB_
index_t KStrideTensorB_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1755
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Do_
index_t Do_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1749
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::InRightPadW_
index_t InRightPadW_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1761
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::CStrideTensorA_
index_t CStrideTensorA_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1755
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::NStrideTensorC_
index_t NStrideTensorC_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1756
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::ConvDilationH_
index_t ConvDilationH_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1759
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::ConvStrideH_
index_t ConvStrideH_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1758
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::ZYX_
index_t ZYX_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1762
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::GStrideTensorA_
index_t GStrideTensorA_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1757
ck::tensor_operation::TransformConvFwdToGemm::AreDescriptorsSmallerThan2GB
__host__ bool AreDescriptorsSmallerThan2GB() const
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:375
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Y_
index_t Y_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1750
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::InLeftPadW_
index_t InLeftPadW_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1760
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::K_
index_t K_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1751
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::GStrideTensorB_
index_t GStrideTensorB_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1757
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::ConvStrideW_
index_t ConvStrideW_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1758
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::C_
index_t C_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1751
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::N_
index_t N_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1747
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::InLeftPadH_
index_t InLeftPadH_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1760
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::DiStride_
index_t DiStride_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1752
ck::tensor_operation::TransformConvFwdToGemm::SplitConvProblem
__host__ auto SplitConvProblem(const ADataType *a_grid_ptr_base, DsPointer &ds_grid_ptr_base, CDataType *c_grid_ptr_base) const
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:393
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::WiStride_
index_t WiStride_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1752
ck::tensor_operation::TransformConvFwdToGemm::MakeADescriptor_M_K
__host__ __device__ auto MakeADescriptor_M_K() const
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:507
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::ConvStrideD_
index_t ConvStrideD_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1758
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Z_
index_t Z_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1750
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::DoStride_
index_t DoStride_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1753
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::HiStride_
index_t HiStride_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1752
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::KStrideTensorC_
index_t KStrideTensorC_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1755
ck::tensor_operation::TransformConvFwdToGemm::TransformConvFwdToGemm
__host__ __device__ TransformConvFwdToGemm(const ConvDimsType &a_g_n_c_wis_lengths, const ConvDimsType &a_g_n_c_wis_strides, const ConvDimsType &b_g_k_c_xs_lengths, const ConvDimsType &b_g_k_c_xs_strides, const ConvDimsType &c_g_n_k_wos_lengths, const ConvDimsType &c_g_n_k_wos_strides, const ConvSpatialDimsType &conv_filter_strides, const ConvSpatialDimsType &conv_filter_dilations, const ConvSpatialDimsType &input_left_pads, const ConvSpatialDimsType &input_right_pads)
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:154
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::X_
index_t X_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1750
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::NStrideTensorA_
index_t NStrideTensorA_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1756
ck::tensor_operation::TransformConvFwdToGemm::MakeCDescriptor_M_N
__host__ __device__ auto MakeCDescriptor_M_N() const
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1425
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Wi_
index_t Wi_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1748
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::ConvDilationD_
index_t ConvDilationD_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1759
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Hi_
index_t Hi_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1748
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Ho_
index_t Ho_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1749
ck::tensor_operation::TransformConvFwdToGemm::TransformConvFwdToGemm
__host__ __device__ TransformConvFwdToGemm(const TransformConvFwdToGemmBase &transform_conv_fwd_to_gemm_base)
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:105
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::HoStride_
index_t HoStride_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1753
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::Wo_
index_t Wo_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1749
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::GStrideTensorC_
index_t GStrideTensorC_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1757
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::WoStride_
index_t WoStride_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1753
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::ConvDilationW_
index_t ConvDilationW_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1759
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::InRightPadD_
index_t InRightPadD_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1761
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::CStrideTensorB_
index_t CStrideTensorB_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1755
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::XStride_
index_t XStride_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1754
ck::tensor_operation::TransformConvFwdToGemm< NDimSpatial, ConvolutionForwardSpecialization::Default >::InLeftPadD_
index_t InLeftPadD_
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1760
ck::tensor_operation::TransformConv::TransformConv
TransformConv()
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1769
ck::tensor_operation::TransformConv::transform_func
auto transform_func(TransformConvFwdToGemm< NDimSpatial, ConvForwardSpecialization > conv_fwd_to_gemm)
Definition tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp:1774
tensor_descriptor.hpp
tensor_descriptor_helper.hpp
tensor_layout.hpp