diff --git a/src/targets/gpu/CMakeLists.txt b/src/targets/gpu/CMakeLists.txt index 2dbaac523ea..5b46eec8b66 100644 --- a/src/targets/gpu/CMakeLists.txt +++ b/src/targets/gpu/CMakeLists.txt @@ -225,6 +225,7 @@ add_library(migraphx_gpu hsa_chiplet.cpp kernel.cpp lower_device_ops.cpp + lower_reshape.cpp lowering.cpp loop.cpp lrn.cpp @@ -237,7 +238,6 @@ add_library(migraphx_gpu prepare_reduce.cpp pooling.cpp problem_cache.cpp - propagate_reshape_layout.cpp rocblas.cpp schedule_model.cpp sync_device.cpp diff --git a/src/targets/gpu/include/migraphx/gpu/lower_reshape.hpp b/src/targets/gpu/include/migraphx/gpu/lower_reshape.hpp new file mode 100644 index 00000000000..975c7ff282b --- /dev/null +++ b/src/targets/gpu/include/migraphx/gpu/lower_reshape.hpp @@ -0,0 +1,89 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 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. + */ +#ifndef MIGRAPHX_GUARD_GPU_LOWER_RESHAPE_HPP +#define MIGRAPHX_GUARD_GPU_LOWER_RESHAPE_HPP + +#include +#include + +namespace migraphx { +inline namespace MIGRAPHX_INLINE_NS { + +struct module; + +namespace gpu { + +/** + * Lower `reshape` into a `reshape_lazy` view, inserting a repacking copy only when the + * input layout does not permit aliasing. + * + * A `reshape` preserves row-major logical element order, so it can be a free view over + * the input buffer whenever the output dims are expressible as a pure restriding of the + * input. Otherwise the bytes have to be moved first. + * + * This pass deliberately runs *after* `eliminate_contiguous`. Lowering inserts a + * `gpu::contiguous` for many ops and `eliminate_contiguous` then drops the redundant + * ones globally, so a reshape's real input layout is not settled until that pass has + * finished. Deciding here means deciding once, with the final layout in hand, instead + * of emitting a pessimistic copy during lowering and undoing it afterwards. + * + * For each `reshape`, in order: + * + * 1. If `reshape_lazy` can alias the input directly and lands on the element count the + * `reshape` declared, emit just the `reshape_lazy`. No copy. The resulting view keeps + * whatever strides aliasing produced, which need not be the ones `reshape` reported. + * + * 2. Otherwise, derive the memory order the input would need in order to be aliasable + * by running the dim mapping backwards (output dims -> input dims), and emit a + * `layout` copy into that order followed by the `reshape_lazy`. The `layout` is + * wrapped in `gpu::precompile_op` and compiles to the existing pointwise copy + * kernel, so this costs the same single kernel as a standardizing contiguous while + * preserving the permutation a following op may want (e.g. NHWC into pooling). + * An identity permutation is declined here and left to case 3, since a `layout` + * that reorders nothing is just a standardizing copy and `gpu::contiguous` has a + * prebuilt code object for it. + * + * 3. Otherwise fall back to a standardizing `gpu::contiguous` followed by the + * `reshape_lazy`. + * + * Cases 1 and 2 need real strides to reason about, so they are skipped for range-based + * dynamic inputs, which have none; those always take case 3. Static inputs are lifted + * to symbolic literals so static and symbolic shapes share one code path. + * + * The 2 input form of `reshape` (`reshape(data, output_buffer)`, where the target shape + * is carried by a runtime-sized output buffer) is not matched and passes through + * un-lowered. No GPU copy op can express it: they all derive their kernel from an index + * space shared by source and destination, which a rank-changing copy does not have. + */ +struct MIGRAPHX_GPU_EXPORT lower_reshape +{ + std::string name() const { return "gpu::lower_reshape"; } + void apply(module& m) const; +}; + +} // namespace gpu +} // namespace MIGRAPHX_INLINE_NS +} // namespace migraphx + +#endif // MIGRAPHX_GUARD_GPU_LOWER_RESHAPE_HPP diff --git a/src/targets/gpu/include/migraphx/gpu/propagate_reshape_layout.hpp b/src/targets/gpu/include/migraphx/gpu/propagate_reshape_layout.hpp deleted file mode 100644 index 4da6b6c589d..00000000000 --- a/src/targets/gpu/include/migraphx/gpu/propagate_reshape_layout.hpp +++ /dev/null @@ -1,65 +0,0 @@ -/* - * The MIT License (MIT) - * - * Copyright (c) 2025-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. - */ -#ifndef MIGRAPHX_GUARD_GPU_PROPAGATE_RESHAPE_LAYOUT_HPP -#define MIGRAPHX_GUARD_GPU_PROPAGATE_RESHAPE_LAYOUT_HPP - -#include -#include - -namespace migraphx { -inline namespace MIGRAPHX_INLINE_NS { - -struct module; - -namespace gpu { - -/** - * Propagate the permutation that a `reshape` would produce through its lowered - * `reshape_lazy`. - * - * Lowering turns `reshape` into `gpu::contiguous -> reshape_lazy -> gpu::contiguous`, - * and `eliminate_contiguous` drops the leading contiguous whenever `reshape_lazy` can - * alias its input directly. When it cannot, the contiguous is kept and forces a - * *standard* layout, discarding the permutation the original `reshape` would have - * propagated (e.g. NHWC into a following pooling op). - * - * This pass runs after `eliminate_contiguous`, where the real input layout is finally - * known. For each surviving `gpu::contiguous -> reshape_lazy` over a non-standard - * input, it replaces the standardizing contiguous with a `layout` that repacks the - * input into the exact memory order `reshape_lazy` needs to alias straight to the - * permuted output. The `layout` reuses the existing pointwise copy kernel, so no - * reshape kernel is generated, and the kernel count is unchanged from the previous - * standardizing contiguous. - */ -struct MIGRAPHX_GPU_EXPORT propagate_reshape_layout -{ - std::string name() const { return "gpu::propagate_reshape_layout"; } - void apply(module& m) const; -}; - -} // namespace gpu -} // namespace MIGRAPHX_INLINE_NS -} // namespace migraphx - -#endif // MIGRAPHX_GUARD_GPU_PROPAGATE_RESHAPE_LAYOUT_HPP diff --git a/src/targets/gpu/lower_reshape.cpp b/src/targets/gpu/lower_reshape.cpp new file mode 100644 index 00000000000..85b3c789bb5 --- /dev/null +++ b/src/targets/gpu/lower_reshape.cpp @@ -0,0 +1,113 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 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 +#include +#include +#include + +namespace migraphx { +inline namespace MIGRAPHX_INLINE_NS { +namespace gpu { + +namespace { +instruction_ref +insert_precompile_op(module& m, instruction_ref pos, instruction_ref input, const operation& op) +{ + auto output_shape = op.compute_shape({input->get_shape()}); + auto alloc = + m.insert_instruction(pos, make_op("allocate", {{"shape", to_value(output_shape)}})); + return m.insert_instruction( + pos, make_op("gpu::precompile_op", {{"op", to_value(op)}}), input, alloc); +} + +instruction_ref insert_contiguous(module& m, instruction_ref pos, instruction_ref input) +{ + const auto& s = input->get_shape(); + shape output_shape = s.dynamic() ? shape{s.type(), s.dyn_dims()} : shape{s.type(), s.lens()}; + auto alloc = + m.insert_instruction(pos, make_op("allocate", {{"shape", to_value(output_shape)}})); + return m.insert_instruction(pos, make_op("gpu::contiguous"), input, alloc); +} + +struct find_reshape : match::supports_dynamic_shapes +{ + // Skip reshape(data, output_buffer). Every GPU copy op derives its kernel from one + // index space shared by source and destination, so none of them can change rank. + auto matcher() const { return match::name("reshape")(match::nargs(1)); } + + void apply(module& m, const match::matcher_result& r) const + { + auto ins = r.result; + auto dims = ins->get_operator().to_value().at("dims"); + auto reshape_op = make_op("reshape_lazy", {{"dims", {dims}}}); + auto input = ins->inputs().front(); + const auto& s = input->get_shape(); + + if(not s.dynamic() or s.symbolic()) + { + auto expected = ins->get_shape().to_symbolic(); + auto output_dims = ins->get_shape().sym_dims(); + auto reshaped = reshape_dims(s.to_symbolic(), output_dims, {.lazy = true}); + if(reshaped and sym::same_symbol(reshaped->sym_elements(), expected.sym_elements())) + { + m.replace_instruction(ins, reshape_op, {input}); + return; + } + + auto relayout = + reshape_dims(ins->get_shape().to_symbolic(), s.sym_dims(), {.lazy = true}); + if(relayout) + { + auto perm = find_permutation(*relayout); + // An identity permutation means the input only needs standardizing, which is + // what gpu::contiguous already does. Prefer it: lower_device_ops turns it into + // a prebuilt code object, whereas a layout under gpu::precompile_op would pay + // for a jit compile to produce the same copy. + if(not std::is_sorted(perm.begin(), perm.end())) + { + auto layout = insert_precompile_op( + m, ins, input, make_op("layout", {{"permutation", perm}})); + m.replace_instruction(ins, reshape_op, {layout}); + return; + } + } + } + + auto contiguous = insert_contiguous(m, ins, input); + m.replace_instruction(ins, reshape_op, {contiguous}); + } +}; +} // namespace + +void lower_reshape::apply(module& m) const { match::find_matches(m, find_reshape{}); } + +} // namespace gpu +} // namespace MIGRAPHX_INLINE_NS +} // namespace migraphx diff --git a/src/targets/gpu/lowering.cpp b/src/targets/gpu/lowering.cpp index a0a6dc67afa..5c803238001 100644 --- a/src/targets/gpu/lowering.cpp +++ b/src/targets/gpu/lowering.cpp @@ -39,9 +39,7 @@ #include #include #include -#include #include -#include #include #include @@ -113,7 +111,6 @@ struct miopen_apply add_nms_op(); add_convolution_backwards_op(); add_select_module_op(); - add_reshape_lazy_op(); add_concat_past_present_op(); add_scan_slice_op(); add_fill_op(); @@ -632,32 +629,6 @@ struct miopen_apply }); } - /** - * Adds reshape lazy to reshape ops that can be aliased instead of copied. - * `gpu::contiguous` are added before and after the reshape; these contiguous - * instructions can be removed by the eliminate_contiguous pass. - */ - void add_reshape_lazy_op() - { - apply_map.emplace("reshape", [=](instruction_ref ins) { - std::vector before_contiguous_args = ins->inputs(); - auto before_alloc = insert_allocation(ins, std::prev(ins)->get_shape()); - before_contiguous_args.push_back(before_alloc); - auto before_contig = - mod->insert_instruction(ins, make_op("gpu::contiguous"), {before_contiguous_args}); - - auto new_reshape_lazy = mod->insert_instruction( - ins, - make_op("reshape_lazy", {{"dims", {ins->get_operator().to_value().at("dims")}}}), - before_contig); - - std::vector after_contiguous_args = {new_reshape_lazy}; - auto after_alloc = insert_allocation(new_reshape_lazy, new_reshape_lazy->get_shape()); - after_contiguous_args.push_back(after_alloc); - return mod->replace_instruction(ins, make_op("gpu::contiguous"), after_contiguous_args); - }); - } - void add_concat_past_present_op() { apply_map.emplace("concat_past_present", [=](instruction_ref ins) { diff --git a/src/targets/gpu/propagate_reshape_layout.cpp b/src/targets/gpu/propagate_reshape_layout.cpp deleted file mode 100644 index 0c2de9699f7..00000000000 --- a/src/targets/gpu/propagate_reshape_layout.cpp +++ /dev/null @@ -1,94 +0,0 @@ -/* - * The MIT License (MIT) - * - * Copyright (c) 2025-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 -#include - -namespace migraphx { -inline namespace MIGRAPHX_INLINE_NS { -namespace gpu { - -namespace { -struct find_reshape_lazy_contiguous : match::supports_dynamic_shapes -{ - // eliminate_contiguous only leaves a standardizing gpu::contiguous in front of a - // reshape_lazy when it could not alias the input directly; that is the only case where a - // permutation was discarded. - auto matcher() const - { - return match::name("reshape_lazy")( - match::arg(0)(match::name("gpu::contiguous").bind("contiguous"))); - } - - void apply(module& m, const match::matcher_result& r) const - { - auto rl = r.result; - auto cont = r.instructions["contiguous"]; - auto input = cont->inputs().front(); - const auto& s = input->get_shape(); - // A standard input has no permutation to propagate; a range-based dynamic input has - // no symbolic dims for reshape_dims/find_permutation to work with. - if(s.standard() or (s.dynamic() and not s.symbolic())) - return; - - auto sym_in = s.to_symbolic(); - - auto permuted = reshape_dims(sym_in, rl->get_shape().sym_dims(), {.lazy = false}); - if(not permuted or permuted->standard()) - return; - // reshape_dims does not check the element count; bail when it provably differs, - // matching reshape_lazy::compute_shape (an indeterminate count is allowed through). - auto out_elems = permuted->sym_elements(); - auto in_elems = sym_in.sym_elements(); - if(sym::strict_less(out_elems, in_elems).value_or(false) or - sym::strict_less(in_elems, out_elems).value_or(false)) - return; - auto relayout = reshape_dims(*permuted, s.sym_dims(), {.lazy = true}); - if(not relayout) - return; - - auto layout_op = make_op("layout", {{"permutation", find_permutation(*relayout)}}); - auto layout_shape = layout_op.compute_shape({s}); - auto alloc = - m.insert_instruction(rl, make_op("allocate", {{"shape", to_value(layout_shape)}})); - auto layout = m.insert_instruction( - rl, make_op("gpu::precompile_op", {{"op", to_value(layout_op)}}), input, alloc); - instruction::replace_argument(rl, cont, layout); - } -}; -} // namespace - -void propagate_reshape_layout::apply(module& m) const -{ - match::find_matches(m, find_reshape_lazy_contiguous{}); -} - -} // namespace gpu -} // namespace MIGRAPHX_INLINE_NS -} // namespace migraphx diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp index c410808fc25..18b98370224 100644 --- a/src/targets/gpu/target.cpp +++ b/src/targets/gpu/target.cpp @@ -78,8 +78,8 @@ #include #include #include +#include #include -#include #include #include #include @@ -230,7 +230,7 @@ struct pipeline_factory lowering{get_context(), options.offload_copy}, eliminate_contiguous{"gpu::contiguous"}, dead_code_elimination{}, - propagate_reshape_layout{}, + lower_reshape{}, dead_code_elimination{}, adjust_allocation{gpu_allocation_model{.use_hip_allocate = false}}, dead_code_elimination{}, diff --git a/test/gpu/adjust_allocation.cpp b/test/gpu/adjust_allocation.cpp index 7e17766509e..784a173431f 100644 --- a/test/gpu/adjust_allocation.cpp +++ b/test/gpu/adjust_allocation.cpp @@ -1,7 +1,7 @@ /* * The MIT License (MIT) * - * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved. + * 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 @@ -23,6 +23,7 @@ */ #include #include +#include #include #include #include @@ -55,6 +56,8 @@ static void run_lowering(migraphx::program& p, bool offload_copy = false) migraphx::dead_code_elimination{}, migraphx::eliminate_contiguous{"gpu::contiguous"}, migraphx::dead_code_elimination{}, + migraphx::gpu::lower_reshape{}, + migraphx::dead_code_elimination{}, migraphx::replace_allocate{migraphx::gpu::gpu_allocation_model{}, offload_copy}, migraphx::dead_code_elimination{}}); } diff --git a/test/gpu/lower_reshape.cpp b/test/gpu/lower_reshape.cpp new file mode 100644 index 00000000000..07f7bf14988 --- /dev/null +++ b/test/gpu/lower_reshape.cpp @@ -0,0 +1,297 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 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 +#include +#include +#include +#include + +static void run_pass(migraphx::module& m) +{ + migraphx::run_passes(m, + {migraphx::eliminate_contiguous{"gpu::contiguous"}, + migraphx::dead_code_elimination{}, + migraphx::gpu::lower_reshape{}, + migraphx::dead_code_elimination{}}); +} + +static migraphx::instruction_ref add_contiguous(migraphx::module& m, + migraphx::instruction_ref input) +{ + const auto& s = input->get_shape(); + auto output_shape = + s.dynamic() ? migraphx::shape{s.type(), s.dyn_dims()} : migraphx::shape{s.type(), s.lens()}; + auto alloc = m.add_instruction( + migraphx::make_op("allocate", {{"shape", migraphx::to_value(output_shape)}})); + return m.add_instruction(migraphx::make_op("gpu::contiguous"), input, alloc); +} + +static migraphx::instruction_ref add_precompile_layout(migraphx::module& m, + migraphx::instruction_ref input, + const std::vector& permutation) +{ + auto op = migraphx::make_op("layout", {{"permutation", permutation}}); + auto alloc = m.add_instruction(migraphx::make_op( + "allocate", {{"shape", migraphx::to_value(op.compute_shape({input->get_shape()}))}})); + return m.add_instruction( + migraphx::make_op("gpu::precompile_op", {{"op", migraphx::to_value(op)}}), input, alloc); +} + +TEST_CASE(lower_standard_reshape) +{ + migraphx::shape input_shape{migraphx::shape::float_type, {2, 3, 4}}; + + migraphx::module m1; + { + auto x = m1.add_parameter("x", input_shape); + auto r = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {6, 4}}}), x); + m1.add_return({r}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", input_shape); + auto r = m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {6, 4}}}), x); + m2.add_return({r}); + } + EXPECT(m1 == m2); +} + +// The 2 input form carries its target shape on the output buffer, which no GPU copy op can +// honor for a rank changing reshape. The matcher does not accept it, so it survives the +// pass untouched rather than being lowered to a copy that reports the input shape. +TEST_CASE(output_buffer_reshape_is_not_lowered) +{ + auto build = [](migraphx::module& m) { + auto x = m.add_parameter("x", {migraphx::shape::float_type, {2, 3, 4}}); + auto output = m.add_parameter("output", {migraphx::shape::float_type, {6, 4}}); + auto r = m.add_instruction(migraphx::make_op("reshape"), x, output); + m.add_return({r}); + }; + + migraphx::module m1; + build(m1); + run_pass(m1); + + migraphx::module m2; + build(m2); + EXPECT(m1 == m2); +} + +TEST_CASE(lower_range_dynamic_reshape) +{ + using dd = migraphx::shape::dynamic_dimension; + migraphx::shape input_shape{migraphx::shape::float_type, {dd{1, 4}, dd{24, 24}}}; + + migraphx::module m1; + { + auto x = m1.add_parameter("x", input_shape); + auto r = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {0, 24}}}), x); + m1.add_return({r}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", input_shape); + auto c = add_contiguous(m2, x); + auto r = m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {0, 24}}}), c); + m2.add_return({r}); + } + EXPECT(m1 == m2); +} + +TEST_CASE(keep_required_contiguous) +{ + migraphx::shape input_shape{migraphx::shape::float_type, {2, 3, 4}, {4, 8, 1}}; + + migraphx::module m1; + { + auto x = m1.add_parameter("x", input_shape); + auto c = add_contiguous(m1, x); + auto r = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {6, 4}}}), c); + m1.add_return({r}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", input_shape); + auto c = add_contiguous(m2, x); + auto r = m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {6, 4}}}), c); + m2.add_return({r}); + } + EXPECT(m1 == m2); +} + +// The input is transposed so case 1 cannot alias it, but the backwards derivation lands on +// the identity permutation. Case 2 declines that and case 3 emits gpu::contiguous, which is +// the same copy without the jit compile a layout would cost. +TEST_CASE(lower_standard_result_with_copy) +{ + migraphx::shape input_shape{migraphx::shape::float_type, {3, 2, 4}, {4, 12, 1}}; + + migraphx::module m1; + { + auto x = m1.add_parameter("x", input_shape); + auto r = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {6, 4}}}), x); + m1.add_return({r}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", input_shape); + auto c = add_contiguous(m2, x); + auto r = m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {6, 4}}}), c); + m2.add_return({r}); + } + EXPECT(m1 == m2); + EXPECT(std::prev(m1.end())->inputs().front()->get_shape().standard()); +} + +TEST_CASE(propagate_reshape_layout) +{ + migraphx::shape input_shape{migraphx::shape::float_type, {1, 1, 1024, 1024}}; + + migraphx::module m1; + { + auto x = m1.add_parameter("x", input_shape); + auto r1 = + m1.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 256, 4, 256, 4}}}), x); + auto t = m1.add_instruction( + migraphx::make_op("transpose", {{"permutation", {0, 2, 4, 1, 3}}}), r1); + auto c = add_contiguous(m1, t); + auto r2 = + m1.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 16, 256, 256}}}), c); + m1.add_return({add_contiguous(m1, r2)}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", input_shape); + auto r1 = m2.add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", {1, 256, 4, 256, 4}}}), x); + auto t = m2.add_instruction( + migraphx::make_op("transpose", {{"permutation", {0, 2, 4, 1, 3}}}), r1); + auto l = add_precompile_layout(m2, t, {0, 3, 4, 1, 2}); + auto r2 = + m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {1, 16, 256, 256}}}), l); + m2.add_return({add_contiguous(m2, r2)}); + } + EXPECT(m1 == m2); + + // Pin the layout the view lands on. A module compare cannot catch a stride change that + // both sides make together, since both derive their strides from the same ops. + migraphx::shape expected_shape{ + migraphx::shape::float_type, {1, 16, 256, 256}, {1048576, 1, 4096, 16}}; + EXPECT(std::prev(m1.end())->inputs().front()->inputs().front()->get_shape() == expected_shape); +} + +// The singleton dims carry arbitrary strides, but the two elements still sit at offsets +// 0 and 1, so dropping a singleton is a pure restriding. eliminate_contiguous drops the +// incoming copy and case 1 aliases the parameter directly; no copy is needed at all. +TEST_CASE(singleton_dims_alias_without_copy) +{ + migraphx::shape input_shape{migraphx::shape::float_type, {1, 1, 2}, {1, 2, 1}}; + + migraphx::module m1; + { + auto x = m1.add_parameter("x", input_shape); + auto c = add_contiguous(m1, x); + auto r = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 2}}}), c); + m1.add_return({add_contiguous(m1, r)}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", input_shape); + auto r = m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {1, 2}}}), x); + m2.add_return({add_contiguous(m2, r)}); + } + EXPECT(m1 == m2); +} + +TEST_CASE(lower_dependent_reshapes) +{ + migraphx::shape input_shape{migraphx::shape::float_type, {1, 4}, {1, 1}}; + + migraphx::module m1; + { + auto x = m1.add_parameter("x", input_shape); + auto c = add_contiguous(m1, x); + auto r1 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 4}}}), c); + auto r2 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2}}}), r1); + m1.add_return({add_contiguous(m1, r2)}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", input_shape); + auto r1 = m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {1, 4}}}), x); + auto r2 = m2.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {2, 2}}}), r1); + m2.add_return({add_contiguous(m2, r2)}); + } + EXPECT(m1 == m2); +} + +// Not converted to a module compare: the expected module would have to name the layout +// permutation, and deriving that by hand from a from_permutation input is exactly the +// thing the test is checking. +TEST_CASE(propagate_symbolic_reshape_layout) +{ + using dd = migraphx::shape::dynamic_dimension; + using migraphx::sym::lit; + + auto n = migraphx::sym::var("N", {1, 8}); + migraphx::shape input_shape = migraphx::shape::from_permutation( + migraphx::shape::float_type, + {dd{n}, dd{lit(4)}, dd{lit(4)}, dd{lit(256)}, dd{lit(256)}}, + {0, 3, 1, 4, 2}); + + migraphx::module m; + auto x = m.add_parameter("x", input_shape); + auto r = m.add_instruction(migraphx::make_op("reshape", {{"dims", {0, 16, 256, 256}}}), x); + m.add_return({r}); + + run_pass(m); + + auto reshape = std::prev(m.end())->inputs().front(); + EXPECT(reshape->name() == "reshape_lazy"); + EXPECT(not reshape->get_shape().standard()); + EXPECT(reshape->get_shape().sym_dims() == + std::vector{n, lit(16), lit(256), lit(256)}); +} + +int main(int argc, const char* argv[]) { test::run(argc, argv); } diff --git a/test/gpu/propagate_reshape_layout.cpp b/test/gpu/propagate_reshape_layout.cpp deleted file mode 100644 index fa238c20495..00000000000 --- a/test/gpu/propagate_reshape_layout.cpp +++ /dev/null @@ -1,184 +0,0 @@ -/* - * The MIT License (MIT) - * - * Copyright (c) 2025-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 -#include -#include "make_precompile_op.hpp" - -static void run_pass(migraphx::module& m) -{ - migraphx::run_passes( - m, {migraphx::gpu::propagate_reshape_layout{}, migraphx::dead_code_elimination{}}); -} - -// A reshape_lazy that cannot alias its non-standard input keeps a standardizing -// gpu::contiguous after eliminate_contiguous. The pass should replace it with a layout -// that repacks the input into the order reshape_lazy needs to propagate the permutation. -TEST_CASE(propagate_permutation) -{ - migraphx::shape in{migraphx::shape::float_type, {1, 1, 1024, 1024}}; - - migraphx::module m1; - { - auto x = m1.add_parameter("x", in); - auto r = - m1.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 256, 4, 256, 4}}}), x); - auto t = m1.add_instruction( - migraphx::make_op("transpose", {{"permutation", {0, 2, 4, 1, 3}}}), r); - // post-eliminate_contiguous state: a standardizing gpu::contiguous feeds reshape_lazy - auto alloc = m1.add_instruction( - migraphx::make_op("allocate", - {{"shape", - migraphx::to_value(migraphx::shape{migraphx::shape::float_type, - t->get_shape().lens()})}})); - auto c = m1.add_instruction(migraphx::make_op("gpu::contiguous"), t, alloc); - auto rl = - m1.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {1, 16, 256, 256}}}), c); - m1.add_return({rl}); - } - run_pass(m1); - - migraphx::module m2; - { - auto x = m2.add_parameter("x", in); - auto r = - m2.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 256, 4, 256, 4}}}), x); - auto t = m2.add_instruction( - migraphx::make_op("transpose", {{"permutation", {0, 2, 4, 1, 3}}}), r); - // layout repacks the transpose into the packed memory order reshape_lazy can alias - auto l_shape = migraphx::shape::from_permutation( - migraphx::shape::float_type, {1, 4, 4, 256, 256}, {0, 3, 4, 1, 2}); - auto alloc = m2.add_instruction( - migraphx::make_op("allocate", {{"shape", migraphx::to_value(l_shape)}})); - auto layout = m2.add_instruction( - make_precompile_op(migraphx::make_op("layout", {{"permutation", {0, 3, 4, 1, 2}}})), - t, - alloc); - auto rl = m2.add_instruction( - migraphx::make_op("reshape_lazy", {{"dims", {1, 16, 256, 256}}}), layout); - m2.add_return({rl}); - } - - EXPECT(m1 == m2); - // reshape_lazy now produces the permuted (NHWC-like) output rather than a standard one - auto rl1 = std::prev(m1.end())->inputs().front(); - EXPECT(rl1->name() == "reshape_lazy"); - EXPECT(not rl1->get_shape().standard()); - EXPECT(rl1->get_shape().lens() == std::vector{1, 16, 256, 256}); -} - -// Symbolic analog of propagate_permutation: the leading batch dimension is a symbol that -// threads through the reshape/transpose/reshape_lazy chain. The pass must propagate the -// permutation just as in the static case, producing the same layout/allocate structure. -TEST_CASE(propagate_permutation_symbolic) -{ - using dd = migraphx::shape::dynamic_dimension; - using migraphx::sym::lit; - - auto n = migraphx::sym::var("N", {1, 8}); - migraphx::shape in{migraphx::shape::float_type, - {dd{n}, dd{lit(1)}, dd{lit(1024)}, dd{lit(1024)}}}; - - migraphx::module m1; - { - auto x = m1.add_parameter("x", in); - // 0 copies the symbolic batch dim; the spatial dims split into blocks. - auto r = - m1.add_instruction(migraphx::make_op("reshape", {{"dims", {0, 256, 4, 256, 4}}}), x); - auto t = m1.add_instruction( - migraphx::make_op("transpose", {{"permutation", {0, 2, 4, 1, 3}}}), r); - // post-eliminate_contiguous state: a standardizing gpu::contiguous feeds reshape_lazy - auto alloc = m1.add_instruction( - migraphx::make_op("allocate", - {{"shape", - migraphx::to_value(migraphx::shape{migraphx::shape::float_type, - t->get_shape().dyn_dims()})}})); - auto c = m1.add_instruction(migraphx::make_op("gpu::contiguous"), t, alloc); - auto rl = - m1.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {0, 16, 256, 256}}}), c); - m1.add_return({rl}); - } - run_pass(m1); - - migraphx::module m2; - { - auto x = m2.add_parameter("x", in); - auto r = - m2.add_instruction(migraphx::make_op("reshape", {{"dims", {0, 256, 4, 256, 4}}}), x); - auto t = m2.add_instruction( - migraphx::make_op("transpose", {{"permutation", {0, 2, 4, 1, 3}}}), r); - // layout repacks the transpose into the packed memory order reshape_lazy can alias - auto l_shape = migraphx::shape::from_permutation( - migraphx::shape::float_type, - {dd{n}, dd{lit(4)}, dd{lit(4)}, dd{lit(256)}, dd{lit(256)}}, - {0, 3, 4, 1, 2}); - auto alloc = m2.add_instruction( - migraphx::make_op("allocate", {{"shape", migraphx::to_value(l_shape)}})); - auto layout = m2.add_instruction( - make_precompile_op(migraphx::make_op("layout", {{"permutation", {0, 3, 4, 1, 2}}})), - t, - alloc); - auto rl = m2.add_instruction( - migraphx::make_op("reshape_lazy", {{"dims", {0, 16, 256, 256}}}), layout); - m2.add_return({rl}); - } - - EXPECT(m1 == m2); - // reshape_lazy now produces the permuted (NHWC-like) symbolic output rather than a standard one - auto rl1 = std::prev(m1.end())->inputs().front(); - EXPECT(rl1->name() == "reshape_lazy"); - EXPECT(not rl1->get_shape().standard()); - EXPECT(rl1->get_shape().sym_dims() == - std::vector{n, lit(16), lit(256), lit(256)}); -} - -// When the reshape collapses the non-standard input back to a standard layout there is no -// permutation to propagate, so the pass must leave the graph unchanged. -TEST_CASE(no_permutation_noop) -{ - migraphx::module m1; - { - auto x = m1.add_parameter("x", {migraphx::shape::float_type, {2, 3, 4}}); - auto t = - m1.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), x); - auto alloc = m1.add_instruction( - migraphx::make_op("allocate", - {{"shape", - migraphx::to_value(migraphx::shape{migraphx::shape::float_type, - t->get_shape().lens()})}})); - auto c = m1.add_instruction(migraphx::make_op("gpu::contiguous"), t, alloc); - auto rl = m1.add_instruction(migraphx::make_op("reshape_lazy", {{"dims", {6, 4}}}), c); - m1.add_return({rl}); - } - migraphx::module m2 = m1; - run_pass(m1); - EXPECT(m1 == m2); -} - -int main(int argc, const char* argv[]) { test::run(argc, argv); }