Improve find_permutation with ambiguous layouts

src/include/migraphx/output_iterator.hpp

@@ -26,7 +26,10 @@
 
 #include <migraphx/config.hpp>
 #include <migraphx/copy_assignable_function.hpp>
+#include <cassert>
 #include <iterator>
+#include <tuple>
+#include <utility>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -65,6 +68,65 @@ function_output_iterator<F> make_function_output_iterator(F f)
     return {std::move(f)};
 }
 
+// Like function_output_iterator, but also advances an underlying iterator so
+// each assignment writes through a different element: `*out = value` calls
+// f(*it, value).
+template <class Iterator, class F>
+struct function_output_iterator_adaptor
+{
+    Iterator it;
+    copy_assignable_function<F> f;
+
+    using self              = function_output_iterator_adaptor;
+    using difference_type   = void;
+    using reference         = void;
+    using value_type        = void;
+    using pointer           = void;
+    using iterator_category = std::output_iterator_tag;
+
+    struct output_proxy
+    {
+        template <class T>
+        output_proxy& operator=(const T& value)
+        {
+            assert(base != nullptr);
+            base->f(*base->it, value);
+            return *this;
+        }
+        self* base;
+    };
+    output_proxy operator*() { return output_proxy{this}; }
+    self& operator++()
+    {
+        ++it;
+        return *this;
+    }
+    self operator++(int) // NOLINT
+    {
+        self result = *this;
+        ++it;
+        return result;
+    }
+};
+
+template <class Iterator, class F>
+function_output_iterator_adaptor<Iterator, F> make_function_output_iterator_adaptor(Iterator it,
+                                                                                    F f)
+{
+    return {std::move(it), std::move(f)};
+}
+
+// Output iterator that assigns through std::get<N> of each element, so
+// algorithms can write to one tuple/pair field of a sequence, such as the
+// values of a map, which an in-place std::transform can't do since the keys are
+// const.
+template <std::size_t N, class Iterator>
+auto element_output_iterator(Iterator it)
+{
+    return make_function_output_iterator_adaptor(
+        std::move(it), [](auto&& x, const auto& value) { std::get<N>(x) = value; });
+}
+
 template <class Container>
 auto join_back_inserter(Container& c)
 {

src/permutation.cpp

@@ -25,6 +25,7 @@
 #include <migraphx/permutation.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/algorithm.hpp>
+#include <migraphx/output_iterator.hpp>
 #include <migraphx/sym.hpp>
 #include <map>
 #include <functional>
@@ -97,23 +98,59 @@ std::vector<int64_t> find_permutation(const shape& s)
     return result;
 }
 
+namespace {
+// A dim of length 1 places no constraint on the memory layout, so a shape
+// supports any permutation that keeps its non-singleton dims in decreasing
+// stride order.
+bool supports_permutation(const shape& s, const std::vector<int64_t>& permutation)
+{
+    assert(permutation.size() == s.ndim());
+    if(s.dynamic())
+        return find_permutation(s) == permutation;
+    std::vector<std::size_t> strides;
+    transform_if(
+        permutation.begin(),
+        permutation.end(),
+        std::back_inserter(strides),
+        [&](auto d) { return s.lens()[d] > 1; },
+        [&](auto d) { return s.strides()[d]; });
+    return std::is_sorted(strides.begin(), strides.end(), std::greater<>{});
+}
+} // namespace
+
 std::vector<int64_t> find_permutation(const std::vector<shape>& shapes)
 {
     if(shapes.empty())
         return {};
-    std::map<std::vector<int64_t>, std::size_t> count;
-    for(auto&& s : shapes)
-    {
-        if(s.broadcasted())
-            continue;
-        count[find_permutation(s)]++;
-    }
-    if(count.empty())
+    std::vector<shape> voters;
+    std::copy_if(shapes.begin(), shapes.end(), std::back_inserter(voters), [](const shape& s) {
+        return not s.broadcasted();
+    });
+    if(voters.empty())
     {
         std::vector<int64_t> r(shapes.front().ndim());
         std::iota(r.begin(), r.end(), 0);
         return r;
     }
+    const auto ndim = voters.front().ndim();
+    if(std::any_of(voters.begin(), voters.end(), [&](const shape& s) { return s.ndim() != ndim; }))
+        MIGRAPHX_THROW("FIND_PERMUTATION: mismatched shape ranks");
+    std::map<std::vector<int64_t>, std::size_t> count;
+    std::transform(voters.begin(),
+                   voters.end(),
+                   std::inserter(count, count.end()),
+                   [](const shape& s) { return std::make_pair(find_permutation(s), 0); });
+    if(count.size() == 1)
+        return count.begin()->first;
+    // When layouts disagree, each shape votes for every candidate it supports.
+    // Shapes with singleton dims are layout-ambiguous and support several, so
+    // they cannot outvote shapes with a definite layout.
+    std::transform(
+        count.begin(), count.end(), element_output_iterator<1>(count.begin()), [&](const auto& p) {
+            return std::count_if(voters.begin(), voters.end(), [&](const shape& s) {
+                return supports_permutation(s, p.first);
+            });
+        });
     auto it = std::max_element(
         count.begin(), count.end(), by(std::less<>{}, [](auto&& p) { return p.second; }));
     assert(it != count.end());

test/layout_convolution.cpp

@@ -362,6 +362,51 @@ TEST_CASE(nhwc_conv_conv)
     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;

test/op_shape_test.cpp

@@ -7310,6 +7310,18 @@ TEST_CASE(test_concat)
     throws_shape(migraphx::make_op("concat", {{"axis", 0}}));
 }
 
+TEST_CASE(test_concat_nhwc_singleton)
+{
+    // The standard-layout input has a singleton channel, so its layout is
+    // ambiguous and the NHWC input decides the output layout.
+    auto sx =
+        migraphx::shape::from_permutation(migraphx::shape::float_type, {1, 47, 8, 8}, {0, 2, 3, 1});
+    migraphx::shape sy{migraphx::shape::float_type, {1, 1, 8, 8}};
+    auto sout =
+        migraphx::shape::from_permutation(migraphx::shape::float_type, {1, 48, 8, 8}, {0, 2, 3, 1});
+    expect_shape(sout, migraphx::make_op("concat", {{"axis", 1}}), sx, sy);
+}
+
 TEST_CASE(test_dyn_concat)
 {
     migraphx::shape sx{migraphx::shape::float_type, {{1, 3, {3}}, {4, 4}, {1, 5, {5}}, {6, 6}}};

test/output_iterator_test.cpp

@@ -0,0 +1,93 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/output_iterator.hpp>
+#include <algorithm>
+#include <map>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "test.hpp"
+
+TEST_CASE(function_output_iterator_collect)
+{
+    std::vector<int> result;
+    std::vector<int> input = {1, 2, 3};
+    std::transform(
+        input.begin(),
+        input.end(),
+        migraphx::make_function_output_iterator([&](const auto& x) { result.push_back(x); }),
+        [](int x) { return x * 2; });
+    EXPECT(result == std::vector<int>{2, 4, 6});
+}
+
+TEST_CASE(join_back_inserter_flatten)
+{
+    std::vector<std::vector<int>> input = {{1, 2}, {}, {3}};
+    std::vector<int> result;
+    std::copy(input.begin(), input.end(), migraphx::join_back_inserter(result));
+    EXPECT(result == std::vector<int>{1, 2, 3});
+}
+
+TEST_CASE(function_output_iterator_adaptor_assign)
+{
+    std::vector<int> result = {0, 0, 0};
+    std::vector<int> input  = {1, 2, 3};
+    std::copy(input.begin(),
+              input.end(),
+              migraphx::make_function_output_iterator_adaptor(
+                  result.begin(), [](int& x, int value) { x = value + 1; }));
+    EXPECT(result == std::vector<int>{2, 3, 4});
+}
+
+TEST_CASE(element_output_iterator_map_values)
+{
+    std::map<std::string, int> m = {{"a", 1}, {"b", 2}, {"c", 3}};
+    std::transform(m.begin(),
+                   m.end(),
+                   migraphx::element_output_iterator<1>(m.begin()),
+                   [](auto&& p) { return p.second * 2; });
+    std::map<std::string, int> expected = {{"a", 2}, {"b", 4}, {"c", 6}};
+    EXPECT(m == expected);
+}
+
+TEST_CASE(element_output_iterator_pair_first)
+{
+    std::vector<std::pair<int, std::string>> v = {{1, "x"}, {2, "y"}};
+    std::vector<int> keys                      = {10, 20};
+    std::copy(keys.begin(), keys.end(), migraphx::element_output_iterator<0>(v.begin()));
+    std::vector<std::pair<int, std::string>> expected = {{10, "x"}, {20, "y"}};
+    EXPECT(v == expected);
+}
+
+TEST_CASE(element_output_iterator_tuple)
+{
+    std::vector<std::tuple<int, int, int>> v = {{1, 2, 3}, {4, 5, 6}};
+    std::vector<int> input                   = {7, 8};
+    std::copy(input.begin(), input.end(), migraphx::element_output_iterator<2>(v.begin()));
+    std::vector<std::tuple<int, int, int>> expected = {{1, 2, 7}, {4, 5, 8}};
+    EXPECT(v == expected);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/shape_test.cpp

@@ -1284,6 +1284,38 @@ TEST_CASE(from_4d_permutation)
     EXPECT(migraphx::find_permutation(out_shape) == permutation);
 }
 
+TEST_CASE(find_permutation_multi_singleton_ambiguous)
+{
+    // A standard shape with a singleton channel is layout-ambiguous, so the
+    // NHWC shape decides the layout.
+    auto nhwc = migraphx::shape::from_permutation(
+        migraphx::shape::float_type, {1, 511, 32, 32}, {0, 2, 3, 1});
+    migraphx::shape single{migraphx::shape::float_type, {1, 1, 32, 32}};
+    std::vector<int64_t> permutation = {0, 2, 3, 1};
+    EXPECT(migraphx::find_permutation({nhwc, single}) == permutation);
+    EXPECT(migraphx::find_permutation({single, nhwc}) == permutation);
+}
+
+TEST_CASE(find_permutation_multi_singleton_only)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {1, 1, 32, 32}};
+    migraphx::shape s2{migraphx::shape::float_type, {1, 1, 32, 32}};
+    std::vector<int64_t> permutation = {0, 1, 2, 3};
+    EXPECT(migraphx::find_permutation({s1, s2}) == permutation);
+}
+
+TEST_CASE(find_permutation_multi_majority)
+{
+    // Shapes without singleton dims keep one vote each, so the majority layout
+    // still wins.
+    auto nhwc =
+        migraphx::shape::from_permutation(migraphx::shape::float_type, {2, 8, 4, 4}, {0, 2, 3, 1});
+    migraphx::shape nchw1{migraphx::shape::float_type, {2, 8, 4, 4}};
+    migraphx::shape nchw2{migraphx::shape::float_type, {2, 8, 4, 4}};
+    std::vector<int64_t> permutation = {0, 1, 2, 3};
+    EXPECT(migraphx::find_permutation({nchw1, nhwc, nchw2}) == permutation);
+}
+
 TEST_CASE(multi_within_bounds)
 {
     migraphx::shape in_shape{migraphx::shape::float_type, {3, 2, 2}};