Skip to content

Commit ed3de88

Browse files
sbalint98sbalint98
committed
Implement Col2Im NHWC layout
1 parent 4e8caf5 commit ed3de88

8 files changed

Lines changed: 390 additions & 41 deletions

File tree

‎projects/miopen/src/gemm_v2.cpp‎

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1788,10 +1788,11 @@ GemmDescriptor CreateGemmDescriptorConvFwd(const conv::ProblemDescription& probl
17881788
}
17891789

17901790
// dx = Col2Im(transpose(w) * dy)
1791-
GemmDescriptor CreateGemmDescriptorConvBwdData(const TensorDescriptor& wDesc,
1792-
const TensorDescriptor& dyDesc,
1793-
const TensorDescriptor& dxDesc)
1791+
GemmDescriptor CreateGemmDescriptorConvBwdData(const conv::ProblemDescription& problem)
17941792
{
1793+
decltype(auto) dyDesc = problem.GetIn();
1794+
decltype(auto) wDesc = problem.GetWeights();
1795+
decltype(auto) dxDesc = problem.GetOut();
17951796
#ifndef NDEBUG
17961797
assert(wDesc.GetType() == dxDesc.GetType() && wDesc.GetType() == dyDesc.GetType());
17971798
#endif
@@ -1804,16 +1805,16 @@ GemmDescriptor CreateGemmDescriptorConvBwdData(const TensorDescriptor& wDesc,
18041805
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
18051806
std::views::take(dyDesc.GetLengths().size() - 2);
18061807

1807-
bool isColMajor = false;
1808-
bool transA = true;
1808+
bool isColMajor = problem.IsLayoutNHWC();
1809+
bool transA = !problem.IsLayoutNHWC();
18091810
bool transB = false;
18101811
int m =
18111812
in_c * std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
18121813
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
18131814
int k = wei_k;
18141815
int lda = m;
1815-
int ldb = n;
1816-
int ldc = n;
1816+
int ldb = problem.IsLayoutNHWC() ? k : n;
1817+
int ldc = problem.IsLayoutNHWC() ? m : n;
18171818
int batch_count = 1;
18181819
auto strideA = static_cast<long long>(0);
18191820
auto strideB = static_cast<long long>(0);
@@ -2215,11 +2216,12 @@ GemmDescriptor CreateGemmDescriptorGroupConvFwd(const conv::ProblemDescription&
22152216
}
22162217

22172218
// dx = Col2Im(transpose(w) * dy)
2218-
GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const TensorDescriptor& wDesc,
2219-
const TensorDescriptor& dyDesc,
2220-
const TensorDescriptor& dxDesc,
2221-
int groupCount)
2219+
GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const conv::ProblemDescription& problem)
22222220
{
2221+
decltype(auto) dyDesc = problem.GetIn();
2222+
decltype(auto) wDesc = problem.GetWeights();
2223+
decltype(auto) dxDesc = problem.GetOut();
2224+
const int groupCount = problem.GetGroupCount();
22232225
#ifndef NDEBUG
22242226
assert(wDesc.GetType() == dxDesc.GetType() && wDesc.GetType() == dyDesc.GetType());
22252227
#endif
@@ -2232,19 +2234,19 @@ GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const TensorDescriptor& wDes
22322234
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
22332235
std::views::take(dyDesc.GetLengths().size() - 2);
22342236

2235-
bool isColMajor = false;
2236-
bool transA = true;
2237+
bool isColMajor = problem.IsLayoutNHWC();
2238+
bool transA = !problem.IsLayoutNHWC();
22372239
bool transB = false;
22382240
int m = (in_c / groupCount) *
22392241
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
22402242
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
22412243
int k = wei_k / groupCount;
22422244
int lda = m;
2243-
int ldb = n;
2244-
int ldc = n;
2245+
int ldb = problem.IsLayoutNHWC() ? groupCount * k : n;
2246+
int ldc = problem.IsLayoutNHWC() ? m : n;
22452247
int batch_count = groupCount;
22462248
auto strideA = static_cast<long long>(m) * k;
2247-
auto strideB = static_cast<long long>(k) * n;
2249+
auto strideB = problem.IsLayoutNHWC() ? k : static_cast<long long>(k) * n;
22482250
auto strideC = static_cast<long long>(m) * n;
22492251
float alpha = 1.;
22502252
float beta = 0.;

‎projects/miopen/src/include/miopen/gemm_v2.hpp‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -182,9 +182,7 @@ GemmDescriptor CreateGemmDescriptorConvFwd(const conv::ProblemDescription& probl
182182

183183
// GEMM parameters for Convolution (using Im2Col) Bwd-Data
184184
// dx = Col2Im(transpose(w) * dy)
185-
GemmDescriptor CreateGemmDescriptorConvBwdData(const TensorDescriptor& wDesc,
186-
const TensorDescriptor& dyDesc,
187-
const TensorDescriptor& dxDesc);
185+
GemmDescriptor CreateGemmDescriptorConvBwdData(const conv::ProblemDescription& problem);
188186

189187
// GEMM parameters for Convolution (using Im2Col) Bwd-Weight
190188
// dw = dy * transpose(Im2Col(x))
@@ -226,10 +224,7 @@ GemmDescriptor CreateGemmDescriptorGroupConvFwd(const conv::ProblemDescription&
226224

227225
// GEMM parameters for Group Convolution (using Im2Col) Bwd-Data
228226
// dx = Col2Im(transpose(w) * dy)
229-
GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const TensorDescriptor& wDesc,
230-
const TensorDescriptor& dyDesc,
231-
const TensorDescriptor& dxDesc,
232-
int groupCount = 1);
227+
GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const conv::ProblemDescription& problem);
233228

234229
// GEMM parameters for Group Convolution (using Im2Col) Bwd-Weight
235230
// dw = dy * transpose(Im2Col(x))

‎projects/miopen/src/include/miopen/util.hpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,9 @@ float Col2ImGPU(const Handle& handle,
4141
const std::vector<size_t>& in_spatial,
4242
Data_t im,
4343
std::size_t im_offset,
44-
miopenDataType_t type);
44+
miopenDataType_t type,
45+
bool layoutNHWC,
46+
const int num_groups);
4547

4648
float Col2Im3dGPUBatched(const Handle& handle,
4749
ConstData_t col,

‎projects/miopen/src/kernels/MIOpenCol2Im2d.cpp‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@
2626
#include "float_types.h"
2727
#include "miopen_cstdint.hpp"
2828

29+
#ifndef LAYOUT_NHWC
30+
#define LAYOUT_NHWC 0
31+
#endif
32+
2933
#if MIOPEN_USE_FP16
3034
#define ACCUMULATOR_NEEDS_CONVERSION 1
3135
#elif MIOPEN_USE_BFP16
@@ -44,6 +48,93 @@ using index_t = uint64_t;
4448
using index_t = uint32_t;
4549
#endif
4650

51+
#if(LAYOUT_NHWC == 1)
52+
extern "C" __global__ void Col2Im2dU(FLOAT* col,
53+
const uint32_t col_h,
54+
const uint32_t col_w,
55+
const uint32_t wei_h,
56+
const uint32_t wei_w,
57+
const uint32_t pad_h,
58+
const uint32_t pad_w,
59+
const uint32_t stride_h,
60+
const uint32_t stride_w,
61+
const uint32_t dilation_h,
62+
const uint32_t dilation_w,
63+
const uint32_t channels,
64+
const uint32_t height,
65+
const uint32_t width,
66+
FLOAT* im,
67+
const uint32_t im_offset)
68+
{
69+
FLOAT* im_off = im + im_offset;
70+
const size_t gid = blockIdx.x * blockDim.x + threadIdx.x;
71+
const size_t global_size = blockDim.x * gridDim.x;
72+
73+
uint32_t c = gid % channels; // coordinates of the image's pixel handled by this thread
74+
const int num_groups = GROUPS;
75+
76+
const int channels_per_group = channels / GROUPS;
77+
const int current_group = c / channels_per_group;
78+
const int channel_in_group = c % channels_per_group;
79+
80+
unsigned int input_size = channels * height * width;
81+
82+
const uint32_t size_of_group = col_h * col_w * (channels / num_groups) * wei_h * wei_w;
83+
84+
uint32_t w = (gid / channels) % width;
85+
uint32_t h = gid / (channels * width);
86+
87+
if(gid >= input_size)
88+
return;
89+
90+
if(h >= height || w >= width || c >= channels)
91+
return;
92+
93+
FLOAT_ACCUM val = (FLOAT_ACCUM)0;
94+
95+
// Loop over all possible (cy, fy) and (cx, fx) such that h = cy + fy and w = cx + fx
96+
// cx,cy - position in the conv output (dy) -- add the filter coordinates (fx,fy) -> you get the
97+
// location in the image, where the filter was applied.
98+
// h + pad_h = cy * stride_h + fy * dilation_h => cy = (h + pad_h - fy * dilation_h) /
99+
// stride_h
100+
for(uint32_t fy = 0; fy < wei_h; fy++)
101+
{
102+
int h_pad = h + pad_h - fy * dilation_h;
103+
if(h_pad < 0 || h_pad % stride_h != 0)
104+
continue;
105+
106+
int cy = h_pad / stride_h;
107+
if(cy < 0 || cy >= col_h)
108+
continue;
109+
110+
for(uint32_t fx = 0; fx < wei_w; fx++)
111+
{
112+
int w_pad = w + pad_w - fx * dilation_w;
113+
if(w_pad < 0 || w_pad % stride_w != 0)
114+
continue;
115+
116+
int cx = w_pad / stride_w;
117+
if(cx < 0 || cx >= col_w)
118+
continue;
119+
120+
size_t col_idx =
121+
(((((cy * col_w + cx) * wei_h + fy) * wei_w + fx) * channels_per_group) +
122+
channel_in_group) +
123+
size_of_group * current_group;
124+
125+
val += CVT_FLOAT2ACCUM(col[col_idx]);
126+
}
127+
}
128+
129+
#if ACCUMULATOR_NEEDS_CONVERSION
130+
im_off[gid] = val > CVT_FLOAT2ACCUM(MAX_VAL) ? MAX_VAL : CVT_ACCUM2FLOAT(val);
131+
#else
132+
im_off[gid] = CVT_ACCUM2FLOAT(val);
133+
#endif
134+
}
135+
136+
#else // LAYOUT_NHWC
137+
47138
extern "C" __global__ void Col2Im2dU(FLOAT* col,
48139
const unsigned int col_h,
49140
const unsigned int col_w,
@@ -106,3 +197,5 @@ extern "C" __global__ void Col2Im2dU(FLOAT* col,
106197
im_off[gid] = tmp;
107198
#endif
108199
}
200+
201+
#endif // LAYOUT_NHWC else

‎projects/miopen/src/kernels/MIOpenCol2Im3d.cpp‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,116 @@
3838
#error "MIOPEN_USE_64BIT_INDEX must be defined"
3939
#endif
4040

41+
#ifndef LAYOUT_NHWC
42+
#define LAYOUT_NHWC 0
43+
#endif
44+
45+
#if(LAYOUT_NHWC == 1)
46+
extern "C" __global__ void Col2Im3dU(FLOAT* col,
47+
const unsigned int col_d,
48+
const unsigned int col_h,
49+
const unsigned int col_w,
50+
const unsigned int wei_d,
51+
const unsigned int wei_h,
52+
const unsigned int wei_w,
53+
const unsigned int pad_d,
54+
const unsigned int pad_h,
55+
const unsigned int pad_w,
56+
const unsigned int stride_d,
57+
const unsigned int stride_h,
58+
const unsigned int stride_w,
59+
const unsigned int dilation_d,
60+
const unsigned int dilation_h,
61+
const unsigned int dilation_w,
62+
const unsigned int channels,
63+
const unsigned int depth,
64+
const unsigned int height,
65+
const unsigned int width,
66+
FLOAT* im,
67+
const uint64_t im_offset)
68+
{
69+
const unsigned int num_groups = GROUPS;
70+
const unsigned int channels_per_group = channels / num_groups;
71+
FLOAT* im_off = im + im_offset;
72+
unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x;
73+
unsigned int global_size = channels * depth * height * width;
74+
if(gid >= global_size)
75+
return;
76+
77+
unsigned int im_ch = gid % channels;
78+
unsigned int group_id = im_ch / channels_per_group;
79+
unsigned int ch_in_group = im_ch % channels_per_group;
80+
81+
unsigned int itmp = gid / channels;
82+
unsigned int im_w = itmp % width;
83+
itmp = itmp / width;
84+
unsigned int im_h = itmp % height;
85+
unsigned int im_d = itmp / height;
86+
87+
im_d += pad_d;
88+
im_h += pad_h;
89+
im_w += pad_w;
90+
91+
unsigned int start_d = (im_d < dilation_d * (wei_d - 1) + 1)
92+
? 0
93+
: (im_d - (dilation_d * (wei_d - 1) + 1)) / stride_d + 1;
94+
unsigned int end_d = min(col_d, im_d / stride_d + 1);
95+
96+
unsigned int start_h = (im_h < dilation_h * (wei_h - 1) + 1)
97+
? 0
98+
: (im_h - (dilation_h * (wei_h - 1) + 1)) / stride_h + 1;
99+
unsigned int end_h = min(col_h, im_h / stride_h + 1);
100+
101+
unsigned int start_w = (im_w < dilation_w * (wei_w - 1) + 1)
102+
? 0
103+
: (im_w - (dilation_w * (wei_w - 1) + 1)) / stride_w + 1;
104+
unsigned int end_w = min(col_w, im_w / stride_w + 1);
105+
106+
uint64_t inner_size = wei_d * wei_h * wei_w * channels_per_group;
107+
uint64_t col_group_size = col_d * col_h * col_w * inner_size;
108+
109+
FLOAT_ACCUM tmp = (FLOAT_ACCUM)0;
110+
111+
for(unsigned int cz = start_d; cz < end_d; cz++)
112+
{
113+
for(unsigned int cy = start_h; cy < end_h; cy++)
114+
{
115+
for(unsigned int cx = start_w; cx < end_w; cx++)
116+
{
117+
if((im_d - cz * stride_d) % dilation_d == 0 &&
118+
(im_h - cy * stride_h) % dilation_h == 0 &&
119+
(im_w - cx * stride_w) % dilation_w == 0)
120+
{
121+
unsigned int z = (im_d - cz * stride_d) / dilation_d;
122+
unsigned int y = (im_h - cy * stride_h) / dilation_h;
123+
unsigned int x = (im_w - cx * stride_w) / dilation_w;
124+
125+
#if MIOPEN_USE_64BIT_INDEX
126+
uint64_t col_off =
127+
group_id * col_group_size +
128+
((((uint64_t)cz * col_h + cy) * col_w + cx) * inner_size) +
129+
(((uint64_t)z * wei_h + y) * wei_w + x) * channels_per_group + ch_in_group;
130+
131+
#else
132+
uint32_t col_off = group_id * col_group_size +
133+
(((cz * col_h + cy) * col_w + cx) * inner_size) +
134+
((z * wei_h + y) * wei_w + x) * channels_per_group +
135+
ch_in_group;
136+
#endif
137+
138+
tmp += CVT_FLOAT2ACCUM(col[col_off]);
139+
}
140+
}
141+
}
142+
}
143+
#if ACCUMULATOR_NEEDS_CONVERSION
144+
im_off[gid] = tmp > CVT_FLOAT2ACCUM(MAX_VAL) ? MAX_VAL : CVT_ACCUM2FLOAT(tmp);
145+
#else
146+
im_off[gid] = tmp;
147+
#endif
148+
}
149+
150+
#else
41151
extern "C" __global__ void Col2Im3dU(FLOAT* col,
42152
const unsigned int col_d,
43153
const unsigned int col_h,
@@ -247,3 +357,4 @@ extern "C" __global__ void Col2Im3dUBatched(FLOAT* col,
247357
im_off[localid] = tmp;
248358
#endif
249359
}
360+
#endif // #if (LAYOUT_NHWC == 1)

0 commit comments

Comments
 (0)