Layout convolution weights as yxck

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 <cstddef>
 #include <string>
+#include <vector>
 #include <migraphx/instruction_ref.hpp>
+#include <migraphx/shape.hpp>
 #include <migraphx/config.hpp>
 
 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<shape::type_t> 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

src/layout_convolution.cpp

@@ -62,6 +62,15 @@ std::vector<int64_t> 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<int64_t> 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<int>() > 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,26 +175,41 @@ 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)
     {
         // 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());
     }
 }
 

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{},

test/layout_convolution.cpp

@@ -32,7 +32,7 @@
 
 #include <test.hpp>
 
 static void run_pass(migraphx::module& m, migraphx::layout_convolution lc = {})
 {
     migraphx::run_passes(m, {lc, migraphx::dead_code_elimination{}});
 }
@@ -52,6 +52,11 @@
     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 @@
     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;