Skip to content

Commit 4e8caf5

Browse files
sbalint98sbalint98
committed
Implement Im2col NHWC layout
1 parent 5cc4f39 commit 4e8caf5

10 files changed

Lines changed: 933 additions & 142 deletions

File tree

‎projects/miopen/src/gemm_v2.cpp‎

Lines changed: 47 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#include <miopen/handle.hpp>
1010
#include <miopen/datatype.hpp>
1111
#include <miopen/hipoc_kernel.hpp>
12+
#include <unistd.h>
1213

1314
#if MIOPEN_USE_HIPBLASLT
1415
#include <hipblaslt/hipblaslt.h>
@@ -704,7 +705,7 @@ miopenStatus_t CallGemm(const Handle& handle,
704705
switch(gemm_desc.dataType)
705706
{
706707
case miopenInt8: {
707-
assert(gemm_desc.k % 4 == 0);
708+
// assert(gemm_desc.k % 4 == 0);
708709

709710
auto alpha = int(gemm_desc.alpha);
710711
auto beta = int(gemm_desc.beta);
@@ -1084,7 +1085,6 @@ miopenStatus_t CallGemmStridedBatched(const Handle& handle,
10841085
switch(gemm_desc.dataType)
10851086
{
10861087
case miopenInt8: {
1087-
assert(gemm_desc.k % 4 == 0);
10881088

10891089
auto alpha = int(gemm_desc.alpha);
10901090
auto beta = int(gemm_desc.beta);
@@ -1732,10 +1732,11 @@ miopenStatus_t CallGemmStridedBatchedSequential(const Handle& handle,
17321732
}
17331733

17341734
// y = w * Im2Col(x)
1735-
GemmDescriptor CreateGemmDescriptorConvFwd(const TensorDescriptor& wDesc,
1736-
const TensorDescriptor& xDesc,
1737-
const TensorDescriptor& yDesc)
1735+
GemmDescriptor CreateGemmDescriptorConvFwd(const conv::ProblemDescription& problem)
17381736
{
1737+
decltype(auto) xDesc = problem.GetIn();
1738+
decltype(auto) wDesc = problem.GetWeights();
1739+
decltype(auto) yDesc = problem.GetOut();
17391740
#ifndef NDEBUG
17401741
assert(wDesc.GetType() == xDesc.GetType());
17411742
if(wDesc.GetType() != miopenInt8)
@@ -1750,16 +1751,16 @@ GemmDescriptor CreateGemmDescriptorConvFwd(const TensorDescriptor& wDesc,
17501751
auto out_spatial =
17511752
yDesc.GetLengths() | std::views::drop(2) | std::views::take(yDesc.GetLengths().size() - 2);
17521753

1753-
bool isColMajor = false;
1754-
bool transA = false;
1755-
bool transB = (wDesc.GetType() == miopenInt8);
1754+
bool isColMajor = problem.IsLayoutNHWC();
1755+
bool transA = problem.IsLayoutNHWC();
1756+
bool transB = problem.IsLayoutNHWC() ? false : (wDesc.GetType() == miopenInt8);
17561757
int m = wei_k;
17571758
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
17581759
int k =
17591760
in_c * std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
17601761
int lda = k;
1761-
int ldb = wDesc.GetType() == miopenInt8 ? k : n;
1762-
int ldc = n;
1762+
int ldb = problem.IsLayoutNHWC() ? k : wDesc.GetType() == miopenInt8 ? k : n;
1763+
int ldc = problem.IsLayoutNHWC() ? m : n;
17631764
int batch_count = 1;
17641765
auto strideA = static_cast<long long>(0);
17651766
auto strideB = static_cast<long long>(0);
@@ -1840,10 +1841,11 @@ GemmDescriptor CreateGemmDescriptorConvBwdData(const TensorDescriptor& wDesc,
18401841
}
18411842

18421843
// dw = dy * transpose(Im2Col(x))
1843-
GemmDescriptor CreateGemmDescriptorConvBwdWeight(const TensorDescriptor& dyDesc,
1844-
const TensorDescriptor& xDesc,
1845-
const TensorDescriptor& dwDesc)
1844+
GemmDescriptor CreateGemmDescriptorConvBwdWeight(const conv::ProblemDescription& problem)
18461845
{
1846+
const auto& dyDesc = problem.GetIn();
1847+
const auto& dwDesc = problem.GetWeights();
1848+
const auto& xDesc = problem.GetOut();
18471849
#ifndef NDEBUG
18481850
assert(dwDesc.GetType() == xDesc.GetType() && dwDesc.GetType() == dyDesc.GetType());
18491851
#endif
@@ -1857,15 +1859,16 @@ GemmDescriptor CreateGemmDescriptorConvBwdWeight(const TensorDescriptor& dyDesc,
18571859
std::views::take(dyDesc.GetLengths().size() - 2);
18581860

18591861
bool isColMajor = false;
1860-
bool transA = false;
1861-
bool transB = true;
1862+
bool transA = problem.IsLayoutNHWC();
1863+
bool transB = !problem.IsLayoutNHWC();
18621864
int m = wei_k;
18631865
int n = static_cast<int>(in_c) *
18641866
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
1865-
int k = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
1866-
int lda = k;
1867-
int ldb = k;
1868-
int ldc = n;
1867+
int k = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
1868+
1869+
int lda = problem.IsLayoutNHWC() ? m : k;
1870+
int ldb = problem.IsLayoutNHWC() ? n : k;
1871+
int ldc = n;
18691872
int batch_count = 1;
18701873
auto strideA = static_cast<long long>(0);
18711874
auto strideB = static_cast<long long>(0);
@@ -2160,14 +2163,12 @@ GemmDescriptor CreateGemmStridedBatchedDescriptorConv1x1BwdWeight(const TensorDe
21602163
}
21612164

21622165
// y = w * Im2Col(x)
2163-
GemmDescriptor CreateGemmDescriptorGroupConvFwd(const TensorDescriptor& wDesc,
2164-
const TensorDescriptor& xDesc,
2165-
const TensorDescriptor& yDesc,
2166-
int groupCount)
2166+
GemmDescriptor CreateGemmDescriptorGroupConvFwd(const conv::ProblemDescription& problem)
21672167
{
2168-
#ifndef NDEBUG
2169-
assert(wDesc.GetType() == xDesc.GetType() && wDesc.GetType() == yDesc.GetType());
2170-
#endif
2168+
decltype(auto) xDesc = problem.GetIn();
2169+
decltype(auto) wDesc = problem.GetWeights();
2170+
decltype(auto) yDesc = problem.GetOut();
2171+
const int groupCount = problem.GetGroupCount();
21712172

21722173
int in_c = xDesc.GetLengths()[1];
21732174
int wei_k = wDesc.GetLengths()[0];
@@ -2177,20 +2178,20 @@ GemmDescriptor CreateGemmDescriptorGroupConvFwd(const TensorDescriptor& wDesc,
21772178
auto out_spatial =
21782179
yDesc.GetLengths() | std::views::drop(2) | std::views::take(yDesc.GetLengths().size() - 2);
21792180

2180-
bool isColMajor = false;
2181-
bool transA = false;
2181+
bool isColMajor = problem.IsLayoutNHWC();
2182+
bool transA = problem.IsLayoutNHWC();
21822183
bool transB = false;
21832184
int m = wei_k / groupCount;
21842185
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
21852186
int k = (in_c / groupCount) *
21862187
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
21872188
int lda = k;
2188-
int ldb = n;
2189-
int ldc = n;
2189+
int ldb = problem.IsLayoutNHWC() ? k : n;
2190+
int ldc = problem.IsLayoutNHWC() ? m * groupCount : n;
21902191
int batch_count = groupCount;
21912192
auto strideA = static_cast<long long>(m) * k;
21922193
auto strideB = static_cast<long long>(k) * n;
2193-
auto strideC = static_cast<long long>(m) * n;
2194+
auto strideC = problem.IsLayoutNHWC() ? m : static_cast<long long>(m) * n;
21942195
float alpha = 1.;
21952196
float beta = 0.;
21962197

@@ -2268,11 +2269,13 @@ GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const TensorDescriptor& wDes
22682269
}
22692270

22702271
// dw = dy * transpose(Im2Col(x))
2271-
GemmDescriptor CreateGemmDescriptorGroupConvBwdWeight(const TensorDescriptor& dyDesc,
2272-
const TensorDescriptor& xDesc,
2273-
const TensorDescriptor& dwDesc,
2274-
int groupCount)
2272+
GemmDescriptor CreateGemmDescriptorGroupConvBwdWeight(const conv::ProblemDescription& problem)
22752273
{
2274+
const auto& dyDesc = problem.GetIn();
2275+
const auto& dwDesc = problem.GetWeights();
2276+
const auto& xDesc = problem.GetOut();
2277+
const auto& conv = problem.GetConv();
2278+
const auto group_count = conv.group_count;
22762279
#ifndef NDEBUG
22772280
assert(dwDesc.GetType() == xDesc.GetType() && dwDesc.GetType() == dyDesc.GetType());
22782281
#endif
@@ -2286,17 +2289,17 @@ GemmDescriptor CreateGemmDescriptorGroupConvBwdWeight(const TensorDescriptor& dy
22862289
std::views::take(dyDesc.GetLengths().size() - 2);
22872290

22882291
bool isColMajor = false;
2289-
bool transA = false;
2290-
bool transB = true;
2291-
int m = wei_k / groupCount;
2292-
int n = (in_c / groupCount) *
2292+
bool transA = problem.IsLayoutNHWC();
2293+
bool transB = !problem.IsLayoutNHWC();
2294+
int m = wei_k / group_count;
2295+
int n = (in_c / group_count) *
22932296
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
22942297
int k = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
2295-
int lda = k;
2296-
int ldb = k;
2297-
int ldc = n;
2298-
int batch_count = groupCount;
2299-
auto strideA = static_cast<long long>(m) * k;
2298+
int lda = problem.IsLayoutNHWC() ? m * group_count : k;
2299+
int ldb = problem.IsLayoutNHWC() ? n : k;
2300+
int ldc = problem.IsLayoutNHWC() ? n : n;
2301+
int batch_count = group_count;
2302+
auto strideA = problem.IsLayoutNHWC() ? m : static_cast<long long>(m) * k;
23002303
auto strideB = static_cast<long long>(k) * n;
23012304
auto strideC = static_cast<long long>(m) * n;
23022305
float alpha = 1.;

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

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
#include <miopen/common.hpp>
3030
#include <miopen/convolution.hpp>
31+
#include <miopen/conv/problem_description.hpp>
3132
#include <miopen/handle.hpp>
3233
#include <miopen/miopen.h>
3334

@@ -177,9 +178,7 @@ CallGemmStridedBatchedSequential(const Handle& handle,
177178

178179
// GEMM parameters for Convolution (using Im2Col) Fwd
179180
// y = w * Im2Col(x)
180-
GemmDescriptor CreateGemmDescriptorConvFwd(const TensorDescriptor& wDesc,
181-
const TensorDescriptor& xDesc,
182-
const TensorDescriptor& yDesc);
181+
GemmDescriptor CreateGemmDescriptorConvFwd(const conv::ProblemDescription& problem);
183182

184183
// GEMM parameters for Convolution (using Im2Col) Bwd-Data
185184
// dx = Col2Im(transpose(w) * dy)
@@ -189,9 +188,7 @@ GemmDescriptor CreateGemmDescriptorConvBwdData(const TensorDescriptor& wDesc,
189188

190189
// GEMM parameters for Convolution (using Im2Col) Bwd-Weight
191190
// dw = dy * transpose(Im2Col(x))
192-
GemmDescriptor CreateGemmDescriptorConvBwdWeight(const TensorDescriptor& dyDesc,
193-
const TensorDescriptor& xDesc,
194-
const TensorDescriptor& dwDesc);
191+
GemmDescriptor CreateGemmDescriptorConvBwdWeight(const conv::ProblemDescription& problem);
195192

196193
// GEMM parameters for 1x1 Convolution (using CNHW) Fwd
197194
// y = CNHW2NCHW(w * NCHW2CNHW(x))
@@ -225,10 +222,7 @@ GemmDescriptor CreateGemmStridedBatchedDescriptorConv1x1BwdWeight(const TensorDe
225222

226223
// GEMM parameters for Group Convolution (using Im2Col) Fwd
227224
// y = w * Im2Col(x)
228-
GemmDescriptor CreateGemmDescriptorGroupConvFwd(const TensorDescriptor& wDesc,
229-
const TensorDescriptor& xDesc,
230-
const TensorDescriptor& yDesc,
231-
int groupCount = 1);
225+
GemmDescriptor CreateGemmDescriptorGroupConvFwd(const conv::ProblemDescription& problem);
232226

233227
// GEMM parameters for Group Convolution (using Im2Col) Bwd-Data
234228
// dx = Col2Im(transpose(w) * dy)
@@ -239,10 +233,7 @@ GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const TensorDescriptor& wDes
239233

240234
// GEMM parameters for Group Convolution (using Im2Col) Bwd-Weight
241235
// dw = dy * transpose(Im2Col(x))
242-
GemmDescriptor CreateGemmDescriptorGroupConvBwdWeight(const TensorDescriptor& dyDesc,
243-
const TensorDescriptor& xDesc,
244-
const TensorDescriptor& dwDesc,
245-
int groupCount = 1);
236+
GemmDescriptor CreateGemmDescriptorGroupConvBwdWeight(const conv::ProblemDescription& problem);
246237

247238
// GEMM parameters for 1x1 Group Convolution (using CNHW) Fwd
248239
// y = CNHW2NCHW(w * NCHW2CNHW(x))

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@ float Im2ColGPU(const Handle& handle,
2525
const std::vector<int>& stride_spatial,
2626
const std::vector<int>& dilation_spatial,
2727
Data_t col,
28-
miopenDataType_t type);
28+
miopenDataType_t type,
29+
bool layoutNHWC,
30+
int num_groups = 1);
2931

3032
float Col2ImGPU(const Handle& handle,
3133
std::size_t spatial_dim,

0 commit comments

Comments
 (0)