Skip to content

Commit a8ba94a

Browse files
orlmon01Copilot
orlmon01
and
Copilot
authored
[MLAS] Add an NHWC implementation of convolution to avoid transposes (microsoft#26834)
* Modification to the CPU EP to specify channels_last when data format is NWHC * Added a FusedNhwcConv kernel * Implementation of the kernel in mlas * Added compiler guards so it is inly used with KleidiAi (for now, can be removed if needed) * Added unittests ### Description Currently OnnxRT supports NCHW as a default datalayout. For optimisations and kernels that operate better in NHWC layout, or where the datalayout is NHWC in the first place Transposes are added around the layers. This patch seeks to eliminate them in cases of convolutions where it would cause a performance decrease. ### Motivation and Context KleidiAi specific implementation of this feature. Only supports convolutions, DepthWise to follow. Currently a little strict with the filters as a result. --------- Signed-off-by: Orlaith Monahan <orlaith.monahan@arm.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 70dcc0e commit a8ba94a

35 files changed

Lines changed: 2036 additions & 779 deletions

‎cmake/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -638,6 +638,7 @@ else()
638638
check_cxx_compiler_flag(-Wcatch-value HAS_CATCH_VALUE)
639639
check_cxx_compiler_flag(-Wclass-memaccess HAS_CLASS_MEMACCESS)
640640
check_cxx_compiler_flag(-Wcharacter-conversion HAS_CHARACTER_CONVERSION)
641+
check_cxx_compiler_flag(-Wno-error=character-conversion HAS_NO_ERROR_CHARACTER_CONVERSION)
641642
check_cxx_compiler_flag(-Wdangling-reference HAS_DANGLING_REFERENCE)
642643
check_cxx_compiler_flag(-Wdeprecated-anon-enum-enum-conversion HAS_DEPRECATED_ANON_ENUM_ENUM_CONVERSION)
643644
check_cxx_compiler_flag(-Wdeprecated-builtins HAS_DEPRECATED_BUILTINS)

‎cmake/onnxruntime_unittests.cmake‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,13 @@ function(filter_test_srcs test_srcs_var)
5050
endfunction()
5151

5252
set(disabled_warnings)
53+
54+
function(onnxruntime_disable_gtest_character_conversion_as_error target_name)
55+
if (HAS_NO_ERROR_CHARACTER_CONVERSION)
56+
target_compile_options(${target_name} PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:-Wno-error=character-conversion>")
57+
endif()
58+
endfunction()
59+
5360
function(AddTest)
5461
cmake_parse_arguments(_UT "DYN" "TARGET" "LIBS;SOURCES;DEPENDS;TEST_ARGS" ${ARGN})
5562
list(REMOVE_DUPLICATES _UT_SOURCES)
@@ -170,9 +177,7 @@ function(AddTest)
170177
if (${HAS_NOERROR})
171178
target_compile_options(${_UT_TARGET} PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:-Wno-error=uninitialized>")
172179
endif()
173-
if (${HAS_CHARACTER_CONVERSION})
174-
target_compile_options(${_UT_TARGET} PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:-Wno-error=character-conversion>")
175-
endif()
180+
onnxruntime_disable_gtest_character_conversion_as_error(${_UT_TARGET})
176181
endif()
177182

178183
set(TEST_ARGS ${_UT_TEST_ARGS})
@@ -847,9 +852,7 @@ if(MSVC)
847852
"$<$<NOT:$<COMPILE_LANGUAGE:CUDA>>:/wd6326>")
848853
else()
849854
target_include_directories(onnxruntime_test_utils PRIVATE ${CMAKE_CURRENT_BINARY_DIR} ${ONNXRUNTIME_ROOT})
850-
if (HAS_CHARACTER_CONVERSION)
851-
target_compile_options(onnxruntime_test_utils PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:-Wno-error=character-conversion>")
852-
endif()
855+
onnxruntime_disable_gtest_character_conversion_as_error(onnxruntime_test_utils)
853856
endif()
854857
if (onnxruntime_USE_NCCL)
855858
target_include_directories(onnxruntime_test_utils PRIVATE ${NCCL_INCLUDE_DIRS})

‎docs/ContribOperators.md‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3569,7 +3569,6 @@ This version of the operator has been available since version 1 of the 'com.micr
35693569
### <a name="com.microsoft.NhwcFusedConv"></a><a name="com.microsoft.nhwcfusedconv">**com.microsoft.NhwcFusedConv**</a>
35703570

35713571
NhwcFusedConv is a Conv operator with optional activation and add operators fused in.
3572-
Only has fp16 implementation as of 2023/04/15.
35733572

35743573
#### Version
35753574

@@ -3600,26 +3599,26 @@ This version of the operator has been available since version 1 of the 'com.micr
36003599

36013600
<dl>
36023601
<dt><tt>X</tt> : T</dt>
3603-
<dd></dd>
3602+
<dd>Input activation tensor in channels-last layout. For 2D convolution this is [N, H, W, C], where N is batch size, H/W are spatial dimensions, and C is the number of input channels.</dd>
36043603
<dt><tt>W</tt> : T</dt>
3605-
<dd></dd>
3604+
<dd>Convolution weight tensor in the standard ONNX Conv filter layout [M, C/group, kH, kW], where M is the number of output channels.</dd>
36063605
<dt><tt>B</tt> (optional) : T</dt>
3607-
<dd></dd>
3606+
<dd>Optional 1D bias tensor of shape [M].</dd>
36083607
<dt><tt>Z</tt> (optional) : T</dt>
3609-
<dd>Tensor to be added to the output, must be the same shape and format as the output tensor.</dd>
3608+
<dd>Optional residual/add tensor in the same channels-last layout and shape as the output tensor Y. For 2D convolution this is [N, out_H, out_W, M].</dd>
36103609
</dl>
36113610

36123611
#### Outputs
36133612

36143613
<dl>
36153614
<dt><tt>Y</tt> : T</dt>
3616-
<dd></dd>
3615+
<dd>Output tensor in channels-last layout. For 2D convolution this is [N, out_H, out_W, M], where M is the number of output channels.</dd>
36173616
</dl>
36183617

36193618
#### Type Constraints
36203619

36213620
<dl>
3622-
<dt><tt>T</tt> : tensor(float16)</dt>
3621+
<dt><tt>T</tt> : tensor(float16), tensor(float)</dt>
36233622
<dd>Constrain input and output types to float tensors</dd>
36243623
</dl>
36253624

‎onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
2020
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization);
2121
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ExpandDims);
2222
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedConv);
23+
#ifdef USE_KLEIDIAI
24+
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NhwcFusedConv);
25+
#endif
2326
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedGemm);
2427
#if !defined(DISABLE_GENERATION_OPS)
2528
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, GreedySearch);
@@ -313,6 +316,9 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
313316
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization)>,
314317
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ExpandDims)>,
315318
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedConv)>,
319+
#ifdef USE_KLEIDIAI
320+
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NhwcFusedConv)>,
321+
#endif
316322
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedGemm)>,
317323
#if !defined(DISABLE_GENERATION_OPS)
318324
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, GreedySearch)>,

‎onnxruntime/contrib_ops/cpu/fused_conv.cc‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,5 +26,18 @@ ONNX_CPU_OPERATOR_TYPED_MS_KERNEL(
2626
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
2727
FusedConvFloat);
2828

29+
#ifdef USE_KLEIDIAI
30+
ONNX_CPU_OPERATOR_TYPED_MS_KERNEL(
31+
NhwcFusedConv,
32+
1,
33+
float,
34+
KernelDefBuilder()
35+
// Allow the optional "sum" input (index 3) to be reused as the output buffer (index 0),
36+
// consistent with the FusedConv kernel registration.
37+
.MayInplace(3, 0)
38+
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
39+
FusedConvFloat);
40+
#endif
41+
2942
} // namespace contrib
3043
} // namespace onnxruntime

0 commit comments

Comments
 (0)