From d4dd53148a17b71cfe35645bad77595f22737e61 Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 18 Aug 2026 11:16:39 -0500 Subject: [PATCH 1/5] Layout weights as yxck --- src/include/migraphx/layout_convolution.hpp | 9 ++ src/layout_convolution.cpp | 41 +++-- src/targets/gpu/target.cpp | 7 +- test/layout_convolution.cpp | 165 ++++++++++++++++++++ 4 files changed, 208 insertions(+), 14 deletions(-) diff --git a/src/include/migraphx/layout_convolution.hpp b/src/include/migraphx/layout_convolution.hpp index 504eb6dc87e..e1357525c68 100644 --- a/src/include/migraphx/layout_convolution.hpp +++ b/src/include/migraphx/layout_convolution.hpp @@ -24,6 +24,7 @@ #ifndef MIGRAPHX_GUARD_MIGRAPHX_LAYOUT_CONVOLUTION_HPP #define MIGRAPHX_GUARD_MIGRAPHX_LAYOUT_CONVOLUTION_HPP +#include #include #include #include @@ -45,6 +46,14 @@ struct MIGRAPHX_EXPORT layout_convolution channels_auto }; layout_order order = channels_first; + // Only used with channels_last: store the weights of fp32 convolutions + // with at least this many output channels with the output channel dim + // innermost (yxck for 2-D convolutions) instead of kyxc. This makes the + // implicit-GEMM A matrix M-contiguous and avoids power-of-2 row strides + // that alias in the memory system. With few output channels there is + // nothing to vectorize along K, so kyxc's dense C loads win. 1 always + // applies the layout; 0 disables it. + std::size_t output_channels_last_threshold = 0; std::string name() const { return "layout_convolution"; } void apply(module_pass_manager& mpm) const; }; diff --git a/src/layout_convolution.cpp b/src/layout_convolution.cpp index 3ae665d1aa9..2aa74bcab5d 100644 --- a/src/layout_convolution.cpp +++ b/src/layout_convolution.cpp @@ -92,7 +92,9 @@ void preserve_output_layout(module& m) } } -void transform_convolutions(module& m, const layout_convolution::layout_order& order) +void transform_convolutions(module& m, + const layout_convolution::layout_order& order, + std::size_t output_channels_last_threshold) { for(auto ins : iterator_for(m)) { @@ -104,11 +106,26 @@ void transform_convolutions(module& m, const layout_convolution::layout_order& o continue; auto v = ins->get_operator().to_value(); bool is_group_conv = v.at("group").to() > 1; + auto perm = is_group_conv ? get_default_permutation(ins) : get_permutation(ins, order); + auto wperm = perm; + // With only a few output channels there is nothing to vectorize along K, + // so keep kyxc where its dense C loads win (e.g. 3-channel RGB heads). + if(output_channels_last_threshold > 0 and order == layout_convolution::channels_last and + not is_group_conv and ins->name() == "convolution" and + ins->inputs().front()->get_shape().type() == shape::float_type and + ins->inputs().back()->get_shape().lens().front() >= output_channels_last_threshold) + { + // Weights [K, C, spatial...] stored spatial-major with the output + // channel dim K innermost (yxck for 2-D convolutions) + std::iota(wperm.begin(), wperm.end() - 2, 2); + *(wperm.end() - 2) = 1; + wperm.back() = 0; + } auto args = ins->inputs(); - auto perm = is_group_conv ? get_default_permutation(ins) : get_permutation(ins, order); - std::transform(args.begin(), args.end(), args.begin(), [&](const auto& i) { - return m.insert_instruction(ins, make_op("layout", {{"permutation", perm}}), i); - }); + args.front() = + m.insert_instruction(ins, make_op("layout", {{"permutation", perm}}), args.front()); + args.back() = + m.insert_instruction(ins, make_op("layout", {{"permutation", wperm}}), args.back()); auto conv = m.insert_instruction(ins, ins->get_operator(), args); auto c = m.insert_instruction(ins, make_op("contiguous"), conv); m.replace_instruction(ins, c); @@ -129,10 +146,12 @@ void remove_layout(module& m) } } -void apply_layout(module& m, layout_convolution::layout_order order) +void apply_layout(module& m, + layout_convolution::layout_order order, + std::size_t output_channels_last_threshold) { preserve_output_layout(m); - transform_convolutions(m, order); + transform_convolutions(m, order, output_channels_last_threshold); run_passes( m, {dead_code_elimination{}, eliminate_contiguous{"contiguous"}, dead_code_elimination{}}); remove_layout(m); @@ -166,18 +185,18 @@ void layout_convolution::apply(module_pass_manager& mpm) const // place with the cheaper one. A copy is not swapped in because its parameters // have fresh identities, which would orphan submodules capturing the originals. module m_first = mpm.get_module(); - apply_layout(m_first, channels_first); + apply_layout(m_first, channels_first, output_channels_last_threshold); module m_last = mpm.get_module(); - apply_layout(m_last, channels_last); + apply_layout(m_last, channels_last, output_channels_last_threshold); // channels_last converts each parameter to NHWC and back, so allow up to two extra // layouts per parameter before preferring channels_first. auto allowance = 2 * mpm.get_module().get_parameters().size(); auto chosen = (score(m_first) + allowance < score(m_last)) ? channels_first : channels_last; - apply_layout(mpm.get_module(), chosen); + apply_layout(mpm.get_module(), chosen, output_channels_last_threshold); } else { - apply_layout(mpm.get_module(), order); + apply_layout(mpm.get_module(), order, output_channels_last_threshold); } } diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp index 8d3bb35667f..47e7f67537e 100644 --- a/src/targets/gpu/target.cpp +++ b/src/targets/gpu/target.cpp @@ -178,9 +178,10 @@ struct pipeline_factory dead_code_elimination{}, rewrite_gelu{options.fast_math}, optimize_module{}, - layout_convolution{.order = enabled(MIGRAPHX_ENABLE_NHWC{}) - ? layout_convolution::channels_last - : layout_convolution::channels_auto}, + layout_convolution{ + .order = enabled(MIGRAPHX_ENABLE_NHWC{}) ? layout_convolution::channels_last + : layout_convolution::channels_auto, + .output_channels_last_threshold = mlir_enabled() ? std::size_t{8} : std::size_t{0}}, dead_code_elimination{}, enable_pass(disabled(MIGRAPHX_ENABLE_FULL_DYNAMIC{}), fuse_horizontal{}), dead_code_elimination{}, diff --git a/test/layout_convolution.cpp b/test/layout_convolution.cpp index 9bc0380d067..80256d06cce 100644 --- a/test/layout_convolution.cpp +++ b/test/layout_convolution.cpp @@ -362,6 +362,171 @@ TEST_CASE(nhwc_conv_conv) EXPECT(m1.sort() == m2.sort()); } +TEST_CASE(nchw_conv_output_channels_last_unchanged) +{ + // output_channels_last only applies to channels_last: NCHW graphs are untouched. + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}}); + auto w = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 8, 3, 3}})); + auto conv = + m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + m1.add_return({conv}); + } + migraphx::module m2 = m1; + run_pass(m1, + {.order = migraphx::layout_convolution::channels_first, + .output_channels_last_threshold = 1}); + EXPECT(m1.sort() == m2.sort()); +} + +TEST_CASE(nhwc_conv_output_channels_last) +{ + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}}); + auto w = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {8, 8, 3, 3}})); + auto conv = + m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + m1.add_return({conv}); + } + run_pass(m1, + {.order = migraphx::layout_convolution::channels_last, + .output_channels_last_threshold = 8}); + + migraphx::module m2; + { + auto x = add_layout_nhwc( + m2, m2.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}})); + auto w = m2.add_instruction(layout({2, 3, 1, 0}), + m2.add_literal(migraphx::generate_literal( + {migraphx::shape::float_type, {8, 8, 3, 3}}))); + auto conv = + m2.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + auto conv_layout = m2.add_instruction(layout(), conv); + m2.add_return({conv_layout}); + } + EXPECT(m1.sort() == m2.sort()); +} + +TEST_CASE(nhwc_conv_output_channels_last_small_k) +{ + // Fewer than 8 output channels keeps kyxc: same result as with the flag off. + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}}); + auto w = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {3, 8, 3, 3}})); + auto conv = + m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + m1.add_return({conv}); + } + migraphx::module m2 = m1; + run_pass(m1, + {.order = migraphx::layout_convolution::channels_last, + .output_channels_last_threshold = 8}); + run_pass(m2, {.order = migraphx::layout_convolution::channels_last}); + EXPECT(m1.sort() == m2.sort()); +} + +TEST_CASE(nhwc_conv_output_channels_last_always) +{ + // A threshold of 1 always stores the weights output-channels-last. + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}}); + auto w = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {3, 8, 3, 3}})); + auto conv = + m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + m1.add_return({conv}); + } + run_pass(m1, + {.order = migraphx::layout_convolution::channels_last, + .output_channels_last_threshold = 1}); + + migraphx::module m2; + { + auto x = add_layout_nhwc( + m2, m2.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}})); + auto w = m2.add_instruction(layout({2, 3, 1, 0}), + m2.add_literal(migraphx::generate_literal( + {migraphx::shape::float_type, {3, 8, 3, 3}}))); + auto conv = + m2.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + auto conv_layout = m2.add_instruction(layout(), conv); + m2.add_return({conv_layout}); + } + EXPECT(m1.sort() == m2.sort()); +} + +TEST_CASE(nhwc_quant_conv_output_channels_last_unchanged) +{ + // yxck applies only to fp32 convolution: quant_convolution keeps kyxc. + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::int8_type, {1, 8, 16, 16}}); + auto w = + m1.add_literal(migraphx::generate_literal({migraphx::shape::int8_type, {4, 8, 3, 3}})); + auto conv = m1.add_instruction( + migraphx::make_op("quant_convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + m1.add_return({conv}); + } + migraphx::module m2 = m1; + run_pass(m1, + {.order = migraphx::layout_convolution::channels_last, + .output_channels_last_threshold = 1}); + run_pass(m2, {.order = migraphx::layout_convolution::channels_last}); + EXPECT(m1.sort() == m2.sort()); +} + +TEST_CASE(nhwc_conv_concat_conv) +{ + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}}); + auto w1 = m1.add_literal( + migraphx::generate_literal({migraphx::shape::float_type, {15, 8, 3, 3}})); + auto conv1 = m1.add_instruction( + migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w1); + auto c = m1.add_literal( + migraphx::generate_literal({migraphx::shape::float_type, {1, 1, 16, 16}})); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), conv1, c); + auto w2 = m1.add_literal( + migraphx::generate_literal({migraphx::shape::float_type, {4, 16, 3, 3}})); + auto conv2 = m1.add_instruction( + migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), concat, w2); + m1.add_return({conv2}); + } + run_pass(m1, {.order = migraphx::layout_convolution::channels_last}); + + migraphx::module m2; + { + auto x = add_layout_nhwc( + m2, m2.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}})); + auto w1 = add_layout_nhwc(m2, + m2.add_literal(migraphx::generate_literal( + {migraphx::shape::float_type, {15, 8, 3, 3}}))); + auto conv1 = m2.add_instruction( + migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w1); + auto c = m2.add_literal( + migraphx::generate_literal({migraphx::shape::float_type, {1, 1, 16, 16}})); + // The singleton-channel literal is layout-ambiguous, so the concat + // output stays channels-last and needs no relayout before conv2. + auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), conv1, c); + auto w2 = add_layout_nhwc(m2, + m2.add_literal(migraphx::generate_literal( + {migraphx::shape::float_type, {4, 16, 3, 3}}))); + auto conv2 = m2.add_instruction( + migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), concat, w2); + auto conv2_layout = m2.add_instruction(layout(), conv2); + m2.add_return({conv2_layout}); + } + EXPECT(m1.sort() == m2.sort()); +} + TEST_CASE(nhwc_conv_reduce) { migraphx::module m1; From b024e12a7e2d7b208278a2bc0884b1783492ebc4 Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 18 Aug 2026 12:34:46 -0500 Subject: [PATCH 2/5] Paramertrize the types --- src/include/migraphx/layout_convolution.hpp | 11 ++-- src/layout_convolution.cpp | 28 +++++----- src/targets/gpu/target.cpp | 5 +- test/layout_convolution.cpp | 57 ++++++++++++++++++++- 4 files changed, 82 insertions(+), 19 deletions(-) diff --git a/src/include/migraphx/layout_convolution.hpp b/src/include/migraphx/layout_convolution.hpp index e1357525c68..b732f5cea53 100644 --- a/src/include/migraphx/layout_convolution.hpp +++ b/src/include/migraphx/layout_convolution.hpp @@ -26,7 +26,9 @@ #include #include +#include #include +#include #include namespace migraphx { @@ -46,14 +48,17 @@ struct MIGRAPHX_EXPORT layout_convolution channels_auto }; layout_order order = channels_first; - // Only used with channels_last: store the weights of fp32 convolutions - // with at least this many output channels with the output channel dim - // innermost (yxck for 2-D convolutions) instead of kyxc. This makes the + // Only used with channels_last: store the weights of convolutions with at + // least this many output channels with the output channel dim innermost + // (yxck for 2-D convolutions) instead of kyxc. This makes the // implicit-GEMM A matrix M-contiguous and avoids power-of-2 row strides // that alias in the memory system. With few output channels there is // nothing to vectorize along K, so kyxc's dense C loads win. 1 always // applies the layout; 0 disables it. std::size_t output_channels_last_threshold = 0; + // Restrict output_channels_last_threshold to weights of these types; when + // empty it applies to all types. + std::vector output_channels_last_types = {}; std::string name() const { return "layout_convolution"; } void apply(module_pass_manager& mpm) const; }; diff --git a/src/layout_convolution.cpp b/src/layout_convolution.cpp index 2aa74bcab5d..1da10137aba 100644 --- a/src/layout_convolution.cpp +++ b/src/layout_convolution.cpp @@ -93,8 +93,8 @@ void preserve_output_layout(module& m) } void transform_convolutions(module& m, - const layout_convolution::layout_order& order, - std::size_t output_channels_last_threshold) + const layout_convolution& options, + layout_convolution::layout_order order) { for(auto ins : iterator_for(m)) { @@ -108,12 +108,14 @@ void transform_convolutions(module& m, bool is_group_conv = v.at("group").to() > 1; auto perm = is_group_conv ? get_default_permutation(ins) : get_permutation(ins, order); auto wperm = perm; + const auto& wshape = ins->inputs().back()->get_shape(); // With only a few output channels there is nothing to vectorize along K, // so keep kyxc where its dense C loads win (e.g. 3-channel RGB heads). - if(output_channels_last_threshold > 0 and order == layout_convolution::channels_last and - not is_group_conv and ins->name() == "convolution" and - ins->inputs().front()->get_shape().type() == shape::float_type and - ins->inputs().back()->get_shape().lens().front() >= output_channels_last_threshold) + if(options.output_channels_last_threshold > 0 and + order == layout_convolution::channels_last and not is_group_conv and + (options.output_channels_last_types.empty() or + contains(options.output_channels_last_types, wshape.type())) and + wshape.lens().front() >= options.output_channels_last_threshold) { // Weights [K, C, spatial...] stored spatial-major with the output // channel dim K innermost (yxck for 2-D convolutions) @@ -147,11 +149,11 @@ void remove_layout(module& m) } void apply_layout(module& m, - layout_convolution::layout_order order, - std::size_t output_channels_last_threshold) + const layout_convolution& options, + layout_convolution::layout_order order) { preserve_output_layout(m); - transform_convolutions(m, order, output_channels_last_threshold); + transform_convolutions(m, options, order); run_passes( m, {dead_code_elimination{}, eliminate_contiguous{"contiguous"}, dead_code_elimination{}}); remove_layout(m); @@ -185,18 +187,18 @@ void layout_convolution::apply(module_pass_manager& mpm) const // place with the cheaper one. A copy is not swapped in because its parameters // have fresh identities, which would orphan submodules capturing the originals. module m_first = mpm.get_module(); - apply_layout(m_first, channels_first, output_channels_last_threshold); + apply_layout(m_first, *this, channels_first); module m_last = mpm.get_module(); - apply_layout(m_last, channels_last, output_channels_last_threshold); + apply_layout(m_last, *this, channels_last); // channels_last converts each parameter to NHWC and back, so allow up to two extra // layouts per parameter before preferring channels_first. auto allowance = 2 * mpm.get_module().get_parameters().size(); auto chosen = (score(m_first) + allowance < score(m_last)) ? channels_first : channels_last; - apply_layout(mpm.get_module(), chosen, output_channels_last_threshold); + apply_layout(mpm.get_module(), *this, chosen); } else { - apply_layout(mpm.get_module(), order, output_channels_last_threshold); + apply_layout(mpm.get_module(), *this, order); } } diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp index 47e7f67537e..31a80a1ef3d 100644 --- a/src/targets/gpu/target.cpp +++ b/src/targets/gpu/target.cpp @@ -181,7 +181,10 @@ struct pipeline_factory layout_convolution{ .order = enabled(MIGRAPHX_ENABLE_NHWC{}) ? layout_convolution::channels_last : layout_convolution::channels_auto, - .output_channels_last_threshold = mlir_enabled() ? std::size_t{8} : std::size_t{0}}, + .output_channels_last_threshold = mlir_enabled() ? std::size_t{8} : std::size_t{0}, + // Only the non-accel path benefits from output-channels-last + // weights; the fp16/int8 accel path prefers kyxc filters. + .output_channels_last_types = {shape::float_type}}, dead_code_elimination{}, enable_pass(disabled(MIGRAPHX_ENABLE_FULL_DYNAMIC{}), fuse_horizontal{}), dead_code_elimination{}, diff --git a/test/layout_convolution.cpp b/test/layout_convolution.cpp index 80256d06cce..c5cfcf4174e 100644 --- a/test/layout_convolution.cpp +++ b/test/layout_convolution.cpp @@ -462,9 +462,61 @@ TEST_CASE(nhwc_conv_output_channels_last_always) EXPECT(m1.sort() == m2.sort()); } +TEST_CASE(nhwc_conv_output_channels_last_all_types) +{ + // An empty type list applies output_channels_last to every type. + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::half_type, {1, 8, 16, 16}}); + auto w = + m1.add_literal(migraphx::generate_literal({migraphx::shape::half_type, {8, 8, 3, 3}})); + auto conv = + m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + m1.add_return({conv}); + } + run_pass(m1, + {.order = migraphx::layout_convolution::channels_last, + .output_channels_last_threshold = 1}); + + migraphx::module m2; + { + auto x = add_layout_nhwc( + m2, m2.add_parameter("x", {migraphx::shape::half_type, {1, 8, 16, 16}})); + auto w = m2.add_instruction( + layout({2, 3, 1, 0}), + m2.add_literal(migraphx::generate_literal({migraphx::shape::half_type, {8, 8, 3, 3}}))); + auto conv = + m2.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + auto conv_layout = m2.add_instruction(layout(), conv); + m2.add_return({conv_layout}); + } + EXPECT(m1.sort() == m2.sort()); +} + +TEST_CASE(nhwc_conv_output_channels_last_type_filtered) +{ + // A type not in output_channels_last_types keeps kyxc. + migraphx::module m1; + { + auto x = m1.add_parameter("x", {migraphx::shape::half_type, {1, 8, 16, 16}}); + auto w = + m1.add_literal(migraphx::generate_literal({migraphx::shape::half_type, {8, 8, 3, 3}})); + auto conv = + m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); + m1.add_return({conv}); + } + migraphx::module m2 = m1; + run_pass(m1, + {.order = migraphx::layout_convolution::channels_last, + .output_channels_last_threshold = 1, + .output_channels_last_types = {migraphx::shape::float_type}}); + run_pass(m2, {.order = migraphx::layout_convolution::channels_last}); + EXPECT(m1.sort() == m2.sort()); +} + TEST_CASE(nhwc_quant_conv_output_channels_last_unchanged) { - // yxck applies only to fp32 convolution: quant_convolution keeps kyxc. + // int8 is not in output_channels_last_types, so quant_convolution keeps kyxc. migraphx::module m1; { auto x = m1.add_parameter("x", {migraphx::shape::int8_type, {1, 8, 16, 16}}); @@ -477,7 +529,8 @@ TEST_CASE(nhwc_quant_conv_output_channels_last_unchanged) migraphx::module m2 = m1; run_pass(m1, {.order = migraphx::layout_convolution::channels_last, - .output_channels_last_threshold = 1}); + .output_channels_last_threshold = 1, + .output_channels_last_types = {migraphx::shape::float_type}}); run_pass(m2, {.order = migraphx::layout_convolution::channels_last}); EXPECT(m1.sort() == m2.sort()); } From 2badd00343a7357e80382978cfa35d4d360e4409 Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 18 Aug 2026 13:14:52 -0500 Subject: [PATCH 3/5] Simplify --- src/include/migraphx/layout_convolution.hpp | 16 +++-- src/layout_convolution.cpp | 60 ++++++++++-------- src/targets/gpu/target.cpp | 16 ++--- test/layout_convolution.cpp | 68 +++++---------------- 4 files changed, 64 insertions(+), 96 deletions(-) diff --git a/src/include/migraphx/layout_convolution.hpp b/src/include/migraphx/layout_convolution.hpp index b732f5cea53..2b618eb8fa8 100644 --- a/src/include/migraphx/layout_convolution.hpp +++ b/src/include/migraphx/layout_convolution.hpp @@ -48,16 +48,14 @@ struct MIGRAPHX_EXPORT layout_convolution channels_auto }; layout_order order = channels_first; - // Only used with channels_last: store the weights of convolutions with at - // least this many output channels with the output channel dim innermost - // (yxck for 2-D convolutions) instead of kyxc. This makes the - // implicit-GEMM A matrix M-contiguous and avoids power-of-2 row strides - // that alias in the memory system. With few output channels there is - // nothing to vectorize along K, so kyxc's dense C loads win. 1 always - // applies the layout; 0 disables it. + // Only used with channels_last: convolutions with at least this many + // output channels store their weights with the K dim innermost (yxck + // instead of kyxc for 2-D); 1 always applies it, 0 disables it. K-innermost + // makes the implicit-GEMM A matrix M-contiguous, avoiding power-of-2 row + // strides. std::size_t output_channels_last_threshold = 0; - // Restrict output_channels_last_threshold to weights of these types; when - // empty it applies to all types. + // Restrict the output-channels-last weight layout to these types; empty + // applies to all types. std::vector output_channels_last_types = {}; std::string name() const { return "layout_convolution"; } void apply(module_pass_manager& mpm) const; diff --git a/src/layout_convolution.cpp b/src/layout_convolution.cpp index 1da10137aba..ec28d4ddde3 100644 --- a/src/layout_convolution.cpp +++ b/src/layout_convolution.cpp @@ -62,6 +62,15 @@ std::vector get_default_permutation(instruction_ref ins) return perm; } +// Weights [K, C, spatial...] stored spatial-major with the output channel dim +// K innermost (yxck for 2-D convolutions) +std::vector get_weight_permutation(instruction_ref ins) +{ + auto perm = get_permutation(ins, layout_convolution::channels_last); + std::rotate(perm.begin(), std::next(perm.begin()), perm.end()); + return perm; +} + bool skip_layout(const shape& s) { return s.ndim() == 1 or s.dynamic() or s.type() == shape::tuple_type; @@ -92,10 +101,10 @@ void preserve_output_layout(module& m) } } -void transform_convolutions(module& m, - const layout_convolution& options, - layout_convolution::layout_order order) +void transform_convolutions(module& m, const layout_convolution& options) { + const bool weights_channels_last = options.output_channels_last_threshold > 0 and + options.order == layout_convolution::channels_last; for(auto ins : iterator_for(m)) { if(not contains({"convolution", "quant_convolution"}, ins->name())) @@ -106,22 +115,21 @@ void transform_convolutions(module& m, continue; auto v = ins->get_operator().to_value(); bool is_group_conv = v.at("group").to() > 1; - auto perm = is_group_conv ? get_default_permutation(ins) : get_permutation(ins, order); + auto perm = + is_group_conv ? get_default_permutation(ins) : get_permutation(ins, options.order); auto wperm = perm; + assert(ins->inputs().size() == 2); const auto& wshape = ins->inputs().back()->get_shape(); - // With only a few output channels there is nothing to vectorize along K, - // so keep kyxc where its dense C loads win (e.g. 3-channel RGB heads). - if(options.output_channels_last_threshold > 0 and - order == layout_convolution::channels_last and not is_group_conv and + // Store channels_last weights K-innermost (yxck) when enabled for this + // weight type and K is large enough to vectorize; below the threshold + // kyxc's dense C loads win. + if(weights_channels_last and not is_group_conv and not wshape.dynamic() and (options.output_channels_last_types.empty() or contains(options.output_channels_last_types, wshape.type())) and wshape.lens().front() >= options.output_channels_last_threshold) { - // Weights [K, C, spatial...] stored spatial-major with the output - // channel dim K innermost (yxck for 2-D convolutions) - std::iota(wperm.begin(), wperm.end() - 2, 2); - *(wperm.end() - 2) = 1; - wperm.back() = 0; + wperm = get_weight_permutation(ins); + assert(wperm.size() == wshape.ndim()); } auto args = ins->inputs(); args.front() = @@ -148,12 +156,12 @@ void remove_layout(module& m) } } -void apply_layout(module& m, - const layout_convolution& options, - layout_convolution::layout_order order) +// Applies options.order, which must be resolved to channels_first or channels_last. +void apply_layout(module& m, const layout_convolution& options) { + assert(options.order != layout_convolution::channels_auto); preserve_output_layout(m); - transform_convolutions(m, options, order); + transform_convolutions(m, options); run_passes( m, {dead_code_elimination{}, eliminate_contiguous{"contiguous"}, dead_code_elimination{}}); remove_layout(m); @@ -181,25 +189,25 @@ std::size_t score(const module& m) void layout_convolution::apply(module_pass_manager& mpm) const { + auto resolved = *this; if(order == layout_order::channels_auto) { // Score each candidate layout on a copy, then transform the live module in // place with the cheaper one. A copy is not swapped in because its parameters // have fresh identities, which would orphan submodules capturing the originals. module m_first = mpm.get_module(); - apply_layout(m_first, *this, channels_first); - module m_last = mpm.get_module(); - apply_layout(m_last, *this, channels_last); + resolved.order = channels_first; + apply_layout(m_first, resolved); + module m_last = mpm.get_module(); + resolved.order = channels_last; + apply_layout(m_last, resolved); // channels_last converts each parameter to NHWC and back, so allow up to two extra // layouts per parameter before preferring channels_first. auto allowance = 2 * mpm.get_module().get_parameters().size(); - auto chosen = (score(m_first) + allowance < score(m_last)) ? channels_first : channels_last; - apply_layout(mpm.get_module(), *this, chosen); - } - else - { - apply_layout(mpm.get_module(), *this, order); + resolved.order = + (score(m_first) + allowance < score(m_last)) ? channels_first : channels_last; } + apply_layout(mpm.get_module(), resolved); } } // namespace MIGRAPHX_INLINE_NS diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp index 31a80a1ef3d..7ca9d430344 100644 --- a/src/targets/gpu/target.cpp +++ b/src/targets/gpu/target.cpp @@ -178,13 +178,15 @@ struct pipeline_factory dead_code_elimination{}, rewrite_gelu{options.fast_math}, optimize_module{}, - layout_convolution{ - .order = enabled(MIGRAPHX_ENABLE_NHWC{}) ? layout_convolution::channels_last - : layout_convolution::channels_auto, - .output_channels_last_threshold = mlir_enabled() ? std::size_t{8} : std::size_t{0}, - // Only the non-accel path benefits from output-channels-last - // weights; the fp16/int8 accel path prefers kyxc filters. - .output_channels_last_types = {shape::float_type}}, + layout_convolution{.order = enabled(MIGRAPHX_ENABLE_NHWC{}) + ? layout_convolution::channels_last + : layout_convolution::channels_auto, + // Below 8 output channels there is too little to vectorize + // along K, so kyxc's dense C loads win (e.g. RGB heads). + .output_channels_last_threshold = mlir_enabled() ? 8u : 0u, + // Only the non-accel path benefits from output-channels-last + // weights; the fp16/int8 accel path prefers kyxc filters. + .output_channels_last_types = {shape::float_type}}, dead_code_elimination{}, enable_pass(disabled(MIGRAPHX_ENABLE_FULL_DYNAMIC{}), fuse_horizontal{}), dead_code_elimination{}, diff --git a/test/layout_convolution.cpp b/test/layout_convolution.cpp index c5cfcf4174e..8915d8959d5 100644 --- a/test/layout_convolution.cpp +++ b/test/layout_convolution.cpp @@ -52,6 +52,11 @@ static migraphx::instruction_ref add_layout_nchw(migraphx::module& m, migraphx:: return m.add_instruction(layout({0, 1, 2, 3}), ins); } +static migraphx::instruction_ref add_layout_yxck(migraphx::module& m, migraphx::instruction_ref ins) +{ + return m.add_instruction(layout({2, 3, 1, 0}), ins); +} + TEST_CASE(auto_conv_nchw) { migraphx::module m1; @@ -400,9 +405,9 @@ TEST_CASE(nhwc_conv_output_channels_last) { auto x = add_layout_nhwc( m2, m2.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}})); - auto w = m2.add_instruction(layout({2, 3, 1, 0}), - m2.add_literal(migraphx::generate_literal( - {migraphx::shape::float_type, {8, 8, 3, 3}}))); + auto w = add_layout_yxck(m2, + m2.add_literal(migraphx::generate_literal( + {migraphx::shape::float_type, {8, 8, 3, 3}}))); auto conv = m2.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); auto conv_layout = m2.add_instruction(layout(), conv); @@ -413,7 +418,7 @@ TEST_CASE(nhwc_conv_output_channels_last) TEST_CASE(nhwc_conv_output_channels_last_small_k) { - // Fewer than 8 output channels keeps kyxc: same result as with the flag off. + // Output channels below the threshold keep kyxc: same result as with the flag off. migraphx::module m1; { auto x = m1.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}}); @@ -451,9 +456,9 @@ TEST_CASE(nhwc_conv_output_channels_last_always) { auto x = add_layout_nhwc( m2, m2.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}})); - auto w = m2.add_instruction(layout({2, 3, 1, 0}), - m2.add_literal(migraphx::generate_literal( - {migraphx::shape::float_type, {3, 8, 3, 3}}))); + auto w = add_layout_yxck(m2, + m2.add_literal(migraphx::generate_literal( + {migraphx::shape::float_type, {3, 8, 3, 3}}))); auto conv = m2.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); auto conv_layout = m2.add_instruction(layout(), conv); @@ -482,8 +487,8 @@ TEST_CASE(nhwc_conv_output_channels_last_all_types) { auto x = add_layout_nhwc( m2, m2.add_parameter("x", {migraphx::shape::half_type, {1, 8, 16, 16}})); - auto w = m2.add_instruction( - layout({2, 3, 1, 0}), + auto w = add_layout_yxck( + m2, m2.add_literal(migraphx::generate_literal({migraphx::shape::half_type, {8, 8, 3, 3}}))); auto conv = m2.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w); @@ -535,51 +540,6 @@ TEST_CASE(nhwc_quant_conv_output_channels_last_unchanged) EXPECT(m1.sort() == m2.sort()); } -TEST_CASE(nhwc_conv_concat_conv) -{ - migraphx::module m1; - { - auto x = m1.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}}); - auto w1 = m1.add_literal( - migraphx::generate_literal({migraphx::shape::float_type, {15, 8, 3, 3}})); - auto conv1 = m1.add_instruction( - migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w1); - auto c = m1.add_literal( - migraphx::generate_literal({migraphx::shape::float_type, {1, 1, 16, 16}})); - auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), conv1, c); - auto w2 = m1.add_literal( - migraphx::generate_literal({migraphx::shape::float_type, {4, 16, 3, 3}})); - auto conv2 = m1.add_instruction( - migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), concat, w2); - m1.add_return({conv2}); - } - run_pass(m1, {.order = migraphx::layout_convolution::channels_last}); - - migraphx::module m2; - { - auto x = add_layout_nhwc( - m2, m2.add_parameter("x", {migraphx::shape::float_type, {1, 8, 16, 16}})); - auto w1 = add_layout_nhwc(m2, - m2.add_literal(migraphx::generate_literal( - {migraphx::shape::float_type, {15, 8, 3, 3}}))); - auto conv1 = m2.add_instruction( - migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), x, w1); - auto c = m2.add_literal( - migraphx::generate_literal({migraphx::shape::float_type, {1, 1, 16, 16}})); - // The singleton-channel literal is layout-ambiguous, so the concat - // output stays channels-last and needs no relayout before conv2. - auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), conv1, c); - auto w2 = add_layout_nhwc(m2, - m2.add_literal(migraphx::generate_literal( - {migraphx::shape::float_type, {4, 16, 3, 3}}))); - auto conv2 = m2.add_instruction( - migraphx::make_op("convolution", {{"padding", {1, 1, 1, 1}}}), concat, w2); - auto conv2_layout = m2.add_instruction(layout(), conv2); - m2.add_return({conv2_layout}); - } - EXPECT(m1.sort() == m2.sort()); -} - TEST_CASE(nhwc_conv_reduce) { migraphx::module m1; From 82d9f32cdb97ba720c7fc136d6419a0280abf41e Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 18 Aug 2026 13:29:40 -0500 Subject: [PATCH 4/5] Cleanup --- src/include/migraphx/layout_convolution.hpp | 3 ++ src/layout_convolution.cpp | 47 +++++++++++---------- 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/src/include/migraphx/layout_convolution.hpp b/src/include/migraphx/layout_convolution.hpp index 2b618eb8fa8..90ffefb5e4d 100644 --- a/src/include/migraphx/layout_convolution.hpp +++ b/src/include/migraphx/layout_convolution.hpp @@ -34,6 +34,7 @@ namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { +struct module; struct module_pass_manager; /** @@ -59,6 +60,8 @@ struct MIGRAPHX_EXPORT layout_convolution std::vector output_channels_last_types = {}; std::string name() const { return "layout_convolution"; } void apply(module_pass_manager& mpm) const; + // Applies this->order, which must be resolved to channels_first or channels_last. + void apply_layout(module& m) const; }; } // namespace MIGRAPHX_INLINE_NS diff --git a/src/layout_convolution.cpp b/src/layout_convolution.cpp index ec28d4ddde3..dc2fab44691 100644 --- a/src/layout_convolution.cpp +++ b/src/layout_convolution.cpp @@ -156,18 +156,6 @@ void remove_layout(module& m) } } -// Applies options.order, which must be resolved to channels_first or channels_last. -void apply_layout(module& m, const layout_convolution& options) -{ - assert(options.order != layout_convolution::channels_auto); - preserve_output_layout(m); - transform_convolutions(m, options); - run_passes( - m, {dead_code_elimination{}, eliminate_contiguous{"contiguous"}, dead_code_elimination{}}); - remove_layout(m); - run_passes(m, {dead_code_elimination{}}); -} - std::size_t score(const module& m) { return std::count_if(m.begin(), m.end(), [](const instruction& ins) { @@ -187,27 +175,42 @@ std::size_t score(const module& m) } } // namespace +void layout_convolution::apply_layout(module& m) const +{ + assert(order != channels_auto); + preserve_output_layout(m); + transform_convolutions(m, *this); + run_passes( + m, {dead_code_elimination{}, eliminate_contiguous{"contiguous"}, dead_code_elimination{}}); + remove_layout(m); + run_passes(m, {dead_code_elimination{}}); +} + void layout_convolution::apply(module_pass_manager& mpm) const { - auto resolved = *this; if(order == layout_order::channels_auto) { // Score each candidate layout on a copy, then transform the live module in // place with the cheaper one. A copy is not swapped in because its parameters // have fresh identities, which would orphan submodules capturing the originals. - module m_first = mpm.get_module(); - resolved.order = channels_first; - apply_layout(m_first, resolved); - module m_last = mpm.get_module(); - resolved.order = channels_last; - apply_layout(m_last, resolved); + layout_convolution first = *this; + first.order = channels_first; + module m_first = mpm.get_module(); + first.apply_layout(m_first); + layout_convolution last = *this; + last.order = channels_last; + module m_last = mpm.get_module(); + last.apply_layout(m_last); // channels_last converts each parameter to NHWC and back, so allow up to two extra // layouts per parameter before preferring channels_first. auto allowance = 2 * mpm.get_module().get_parameters().size(); - resolved.order = - (score(m_first) + allowance < score(m_last)) ? channels_first : channels_last; + const auto& chosen = (score(m_first) + allowance < score(m_last)) ? first : last; + chosen.apply_layout(mpm.get_module()); + } + else + { + apply_layout(mpm.get_module()); } - apply_layout(mpm.get_module(), resolved); } } // namespace MIGRAPHX_INLINE_NS From 40378eb945ccfda3d78b3e14460761021a1f6a5e Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 18 Aug 2026 13:29:43 -0500 Subject: [PATCH 5/5] Format --- src/layout_convolution.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/layout_convolution.cpp b/src/layout_convolution.cpp index dc2fab44691..4fd428584ee 100644 --- a/src/layout_convolution.cpp +++ b/src/layout_convolution.cpp @@ -204,7 +204,7 @@ void layout_convolution::apply(module_pass_manager& mpm) const // channels_last converts each parameter to NHWC and back, so allow up to two extra // layouts per parameter before preferring channels_first. auto allowance = 2 * mpm.get_module().get_parameters().size(); - const auto& chosen = (score(m_first) + allowance < score(m_last)) ? first : last; + const auto& chosen = (score(m_first) + allowance < score(m_last)) ? first : last; chosen.apply_layout(mpm.get_module()); } else