Slice over dynamic dimension

CHANGELOG.md

@@ -51,6 +51,7 @@ Full documentation for MIGraphX is available at
 * Updated `QLinearConv` bias handling to dequantize bias using the product of input and weight scales before adding to the convolution output.
 * Updated netron output to create an ONNX-like protobuff. Now also includes debug symbols if enabled. (#4701)
 * Updated python API to allow getting and adding debug symbols from instructions. (#4803)
+* Allow for 1 arg slicing over a dynamic dimension. (#5015)
 
 ### Resolved issues
 

src/include/migraphx/op/slice.hpp

@@ -277,7 +277,20 @@ struct slice
                return not input_shape.dyn_dims()[axis].is_fixed();
            }))
         {
-            MIGRAPHX_THROW("SLICE 1_arg: slicing is not allowed on non-fixed dynamic input axis ");
+            if(input_shape.symbolic())
+            {
+                MIGRAPHX_THROW(
+                    "SLICE 1_arg: slicing is not allowed on non-fixed symbolic input axis ");
+            }
+            // Attributes are not normalized for this case, so they can be negative or
+            // out-of-bounds. Using a relaxed dimension bound for now instead of calculating the
+            // tightest possible bound.
+            auto dds = input_shape.dyn_dims();
+            for(auto axis : this->axes)
+            {
+                dds[axis] = {0, dds[axis].get_interval().max};
+            }
+            return shape{input_shape.type(), dds};
         }
 
         auto new_lens = lens_calc(input_shape.max_lens(), this->starts, this->ends, this->axes);

test/onnx/parse/slice_step_dyn_test.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 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
@@ -26,9 +26,28 @@
 
 TEST_CASE(slice_step_dyn_test)
 {
-    // A slice command with non-default steps will have a "Step" instruction added in parsing.
-    // At the time of writing, Step doesn't support dynamic shape input.
+    // A slice command with non-default steps will have a "step" instruction added in parsing.
+    // Slicing over a non-fixed dynamic dimension and the step op both support dynamic shapes,
+    // so parsing succeeds.
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto l0 =
+        mm->add_parameter("0", migraphx::shape{migraphx::shape::float_type, {{1, 4}, {5, 5}}});
+    // literals from parser
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {2, 1}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, -2}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, -1}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-5, -3}});
+    auto slice_out = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {-1, -2}}, {"starts", {-5, -3}}, {"ends", {-1, -1}}}),
+        l0);
+    auto step_out = mm->add_instruction(
+        migraphx::make_op("step", {{"axes", {-1, -2}}, {"steps", {2, 1}}}), slice_out);
+    mm->add_return({step_out});
+
     migraphx::onnx_options options;
     options.default_dyn_dim_value = {1, 4};
-    EXPECT(test::throws([&] { read_onnx("slice_step_dyn_test.onnx", options); }));
+    auto prog                     = read_onnx("slice_step_dyn_test.onnx", options);
+
+    EXPECT(p == prog);
 }

test/op_shape_test.cpp

@@ -5217,16 +5217,12 @@ TEST_CASE(slice_dyn_shape2)
 
 TEST_CASE(slice_dyn_shape3)
 {
-    // TODO: When non-fixed dimension slicing is allowed, Slice to a size smaller than min.
-    // Until then, this action is an error.
+    // Slicing a non-fixed dynamic dimension relaxes only the sliced axis to {0, max};
+    // non-sliced axes are left unchanged.
     migraphx::shape input{migraphx::shape::int32_type, {{2, 3}, {7, 8}, {2, 3}}};
-    throws_shape(migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}),
+    expect_shape(migraphx::shape{migraphx::shape::int32_type, {{2, 3}, {0, 8}, {2, 3}}},
+                 migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}),
                  input);
-    // clang-format off
-    //     expect_shape(migraphx::shape{migraphx::shape::int32_type, {{2, 3}, {1, 1}, {2, 3}}},
-    //                  migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}),
-    //                  input);
-    // clang-format on
 }
 
 TEST_CASE(slice_dyn_shape4)
@@ -5258,6 +5254,47 @@ TEST_CASE(slice_dyn_preserves_optimals)
                  input);
 }
 
+TEST_CASE(slice_dyn_nonfixed_axis)
+{
+    // Slicing a single non-fixed dynamic axis relaxes only that axis to {0, max}.
+    migraphx::shape input{migraphx::shape::int32_type, {{2, 3}, {7, 8}, {2, 3}}};
+    expect_shape(migraphx::shape{migraphx::shape::int32_type, {{2, 3}, {0, 8}, {2, 3}}},
+                 migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {4}}}),
+                 input);
+}
+
+TEST_CASE(slice_dyn_nonfixed_multi_axis)
+{
+    // Slicing multiple axes where one is non-fixed: only the sliced axes are relaxed,
+    // the non-sliced axis (2) is left unchanged.
+    migraphx::shape input{migraphx::shape::int32_type, {{2, 4}, {7, 8}, {2, 3}}};
+    expect_shape(
+        migraphx::shape{migraphx::shape::int32_type, {{0, 4}, {0, 8}, {2, 3}}},
+        migraphx::make_op("slice", {{"axes", {0, 1}}, {"starts", {1, 1}}, {"ends", {2, 4}}}),
+        input);
+}
+
+TEST_CASE(slice_dyn_fixed_axis_with_nonfixed_neighbors)
+{
+    // Only a fixed axis (1) is sliced; the non-fixed neighbors (0, 2) are not sliced,
+    // so the relax branch is not taken and the slice computes normally.
+    migraphx::shape input{migraphx::shape::int32_type, {{2, 4}, {7, 7}, {2, 5}}};
+    expect_shape(migraphx::shape{migraphx::shape::int32_type, {{2, 4}, {3, 3}, {2, 5}}},
+                 migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {4}}}),
+                 input);
+}
+
+TEST_CASE(slice_dyn_nonfixed_keeps_other_optimals)
+{
+    // Non-sliced axes keep their optimals; the sliced non-fixed axis is relaxed to {0, max}.
+    migraphx::shape input{migraphx::shape::int32_type, {dd{2, 4, {3}}, dd{7, 8}, dd{2, 5, {3, 4}}}};
+    migraphx::shape expected{migraphx::shape::int32_type,
+                             {dd{2, 4, {3}}, dd{0, 8}, dd{2, 5, {3, 4}}}};
+    expect_shape(expected,
+                 migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {4}}}),
+                 input);
+}
+
 TEST_CASE(slice_sym)
 {
     auto n                                      = var("n", {1, 8});