diff --git a/src/include/migraphx/output_iterator.hpp b/src/include/migraphx/output_iterator.hpp index b7c298ceaec..cc49a59e533 100644 --- a/src/include/migraphx/output_iterator.hpp +++ b/src/include/migraphx/output_iterator.hpp @@ -26,7 +26,10 @@ #include #include +#include #include +#include +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { @@ -65,6 +68,65 @@ function_output_iterator 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 +struct function_output_iterator_adaptor +{ + Iterator it; + copy_assignable_function 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 + 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 +function_output_iterator_adaptor make_function_output_iterator_adaptor(Iterator it, + F f) +{ + return {std::move(it), std::move(f)}; +} + +// Output iterator that assigns through std::get 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 +auto element_output_iterator(Iterator it) +{ + return make_function_output_iterator_adaptor( + std::move(it), [](auto&& x, const auto& value) { std::get(x) = value; }); +} + template auto join_back_inserter(Container& c) { diff --git a/src/permutation.cpp b/src/permutation.cpp index b8652617b7d..2f496d62df3 100644 --- a/src/permutation.cpp +++ b/src/permutation.cpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -97,23 +98,59 @@ std::vector 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& permutation) +{ + assert(permutation.size() == s.ndim()); + if(s.dynamic()) + return find_permutation(s) == permutation; + std::vector 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 find_permutation(const std::vector& shapes) { if(shapes.empty()) return {}; - std::map, std::size_t> count; - for(auto&& s : shapes) - { - if(s.broadcasted()) - continue; - count[find_permutation(s)]++; - } - if(count.empty()) + std::vector voters; + std::copy_if(shapes.begin(), shapes.end(), std::back_inserter(voters), [](const shape& s) { + return not s.broadcasted(); + }); + if(voters.empty()) { std::vector 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::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()); diff --git a/test/layout_convolution.cpp b/test/layout_convolution.cpp index 9bc0380d067..b7030786de7 100644 --- a/test/layout_convolution.cpp +++ b/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; diff --git a/test/op_shape_test.cpp b/test/op_shape_test.cpp index f641a5860fd..b2f862b4cff 100644 --- a/test/op_shape_test.cpp +++ b/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}}}; diff --git a/test/output_iterator_test.cpp b/test/output_iterator_test.cpp new file mode 100644 index 00000000000..2cbc8764b1d --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +#include "test.hpp" + +TEST_CASE(function_output_iterator_collect) +{ + std::vector result; + std::vector 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{2, 4, 6}); +} + +TEST_CASE(join_back_inserter_flatten) +{ + std::vector> input = {{1, 2}, {}, {3}}; + std::vector result; + std::copy(input.begin(), input.end(), migraphx::join_back_inserter(result)); + EXPECT(result == std::vector{1, 2, 3}); +} + +TEST_CASE(function_output_iterator_adaptor_assign) +{ + std::vector result = {0, 0, 0}; + std::vector 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{2, 3, 4}); +} + +TEST_CASE(element_output_iterator_map_values) +{ + std::map 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 expected = {{"a", 2}, {"b", 4}, {"c", 6}}; + EXPECT(m == expected); +} + +TEST_CASE(element_output_iterator_pair_first) +{ + std::vector> v = {{1, "x"}, {2, "y"}}; + std::vector keys = {10, 20}; + std::copy(keys.begin(), keys.end(), migraphx::element_output_iterator<0>(v.begin())); + std::vector> expected = {{10, "x"}, {20, "y"}}; + EXPECT(v == expected); +} + +TEST_CASE(element_output_iterator_tuple) +{ + std::vector> v = {{1, 2, 3}, {4, 5, 6}}; + std::vector input = {7, 8}; + std::copy(input.begin(), input.end(), migraphx::element_output_iterator<2>(v.begin())); + std::vector> expected = {{1, 2, 7}, {4, 5, 8}}; + EXPECT(v == expected); +} + +int main(int argc, const char* argv[]) { test::run(argc, argv); } diff --git a/test/shape_test.cpp b/test/shape_test.cpp index 51d7eccf9fb..286e95c1b4b 100644 --- a/test/shape_test.cpp +++ b/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 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 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 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}};