diff --git a/src/include/migraphx/layout_convolution.hpp b/src/include/migraphx/layout_convolution.hpp index 504eb6dc87e..90ffefb5e4d 100644 --- a/src/include/migraphx/layout_convolution.hpp +++ b/src/include/migraphx/layout_convolution.hpp @@ -24,13 +24,17 @@ #ifndef MIGRAPHX_GUARD_MIGRAPHX_LAYOUT_CONVOLUTION_HPP #define MIGRAPHX_GUARD_MIGRAPHX_LAYOUT_CONVOLUTION_HPP +#include #include +#include #include +#include #include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { +struct module; struct module_pass_manager; /** @@ -45,8 +49,19 @@ struct MIGRAPHX_EXPORT layout_convolution channels_auto }; layout_order order = channels_first; + // 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 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; + // 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 3ae665d1aa9..4fd428584ee 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,8 +101,10 @@ 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& 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())) @@ -104,11 +115,27 @@ 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, options.order); + auto wperm = perm; + assert(ins->inputs().size() == 2); + const auto& wshape = ins->inputs().back()->get_shape(); + // 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) + { + wperm = get_weight_permutation(ins); + assert(wperm.size() == wshape.ndim()); + } 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,16 +156,6 @@ void remove_layout(module& m) } } -void apply_layout(module& m, layout_convolution::layout_order order) -{ - preserve_output_layout(m); - transform_convolutions(m, order); - 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) { @@ -158,6 +175,17 @@ 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 { if(order == layout_order::channels_auto) @@ -165,19 +193,23 @@ void layout_convolution::apply(module_pass_manager& mpm) const // 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, channels_first); - module m_last = mpm.get_module(); - apply_layout(m_last, channels_last); + 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(); - auto chosen = (score(m_first) + allowance < score(m_last)) ? channels_first : channels_last; - apply_layout(mpm.get_module(), chosen); + const auto& chosen = (score(m_first) + allowance < score(m_last)) ? first : last; + chosen.apply_layout(mpm.get_module()); } else { - apply_layout(mpm.get_module(), order); + apply_layout(mpm.get_module()); } } diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp index 8d3bb35667f..7ca9d430344 100644 --- a/src/targets/gpu/target.cpp +++ b/src/targets/gpu/target.cpp @@ -180,7 +180,13 @@ struct pipeline_factory optimize_module{}, layout_convolution{.order = enabled(MIGRAPHX_ENABLE_NHWC{}) ? layout_convolution::channels_last - : layout_convolution::channels_auto}, + : 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 9bc0380d067..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; @@ -362,6 +367,179 @@ 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 = 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); + m2.add_return({conv_layout}); + } + EXPECT(m1.sort() == m2.sort()); +} + +TEST_CASE(nhwc_conv_output_channels_last_small_k) +{ + // 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}}); + 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 = 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); + m2.add_return({conv_layout}); + } + 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 = 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); + 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) +{ + // 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}}); + 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, + .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_conv_reduce) { migraphx::module m1;