diff --git a/onnxscript/rewriter/rules/common/_collapse_slices.py b/onnxscript/rewriter/rules/common/_collapse_slices.py index 712aa410f3..7f45997de3 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices.py @@ -55,13 +55,21 @@ def _check_if_redundant_slice( if starts_const.numpy().item() != 0: logger.info("The value 'start' is not 0.") return False + + axis = int(axes_const.numpy().item()) + if data.shape is not None: + rank = data.shape.rank() + if axis < -rank or axis >= rank: + logger.info("The value 'axis' is out of range for the data rank.") + return False + # In case data.shape is not statically known, we still can tell the slice is redundant if ends is sys.maxsize if ends_const.numpy().item() == _INT64_MAX: return True - if data.shape is None or data.shape.is_dynamic(axes_const.numpy().item()): + if data.shape is None or data.shape.is_dynamic(axis): logger.info("The value 'data' shape is not statically known.") return False - if ends_const.numpy().item() < data.shape[axes_const.numpy().item()]: + if ends_const.numpy().item() < data.shape[axis]: logger.info("The value 'end' is less than the shape of the specified axis.") return False @@ -78,7 +86,15 @@ def _potential_redundant_slice(op, data, starts, ends, axes, steps): return op.Slice(data, starts, ends, axes, steps, _outputs=["slice_output"]) -def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_): +def _same_shape( + op, + data: ir.Value, + slice_output: ir.Value, + steps: ir.Value, + starts: ir.Value | None = None, + axes: ir.Value | None = None, + **_, +) -> bool: """Check if the shape of the slice output is the same as the data.""" if data.shape is None or slice_output.shape is None: return False @@ -91,6 +107,41 @@ def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_ elif not _ir_utils.is_singleton_value(steps, 1): return False + rank = data.shape.rank() + + axes_np = _ir_utils.get_numpy_value(axes) if axes is not None else None + if axes_np is not None: + for a in axes_np.flat: + if a < -rank or a >= rank: + return False + + if starts is not None: + starts_np = _ir_utils.get_numpy_value(starts) + if starts_np is not None: + if axes_np is not None: + axes_list = [int(a) for a in axes_np.flat] + if len(axes_list) != len(starts_np.flat): + return False + elif axes is None: + axes_list = list(range(len(starts_np.flat))) + else: + if not all(s == 0 for s in starts_np.flat): + return False + axes_list = [] + + for s, a in zip(starts_np.flat, axes_list): + if a < -rank or a >= rank: + return False + if data.shape.is_dynamic(a): + if s != 0: + return False + else: + dim = data.shape[a] + eff_start = s + dim if s < 0 else s + eff_start = max(0, min(dim, eff_start)) + if eff_start != 0: + return False + return _ir_utils.same_shape(data.shape, slice_output.shape) diff --git a/onnxscript/rewriter/rules/common/_collapse_slices_test.py b/onnxscript/rewriter/rules/common/_collapse_slices_test.py index 2c8101a875..520bab1228 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices_test.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices_test.py @@ -178,3 +178,183 @@ def test_multi_element_steps_numerical_correctness(self): model, (np.random.rand(4, 5).astype(np.float32),), ) + + def test_slice_out_of_bounds_axis_does_not_crash_or_collapse(self): + """Slice with out-of-bounds axis should be rejected without IndexError.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_negative_out_of_bounds_axis_does_not_crash_or_collapse(self): + """Slice with negative out-of-bounds axis should be rejected without IndexError.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_out_of_bounds_axis_with_int64_max_ends_does_not_collapse(self): + """Slice with out-of-bounds axis must not collapse even if ends reaches int64_max.""" + model = ir.from_onnx_text( + f""" + + agraph (float[4, 8] data) => (float[4, 8] output) + {{ + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + }} + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_valid_negative_axis_collapses(self): + """Slice with a valid negative axis (e.g. -1 for rank 2) should collapse.""" + model_text = """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + original = ir.from_onnx_text(model_text) + model = ir.from_onnx_text(model_text) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertIn("Identity", [node.op_type for node in model.graph]) + testing.assert_numerically_equal( + original, + model, + (np.random.rand(4, 8).astype(np.float32),), + ) + + def test_slice_negative_normalized_start_collapses(self): + """Slice with negative start normalizing to 0 (e.g., -8 on dim 8) should collapse.""" + model_text = """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + original = ir.from_onnx_text(model_text) + model = ir.from_onnx_text(model_text) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertIn("Identity", [node.op_type for node in model.graph]) + testing.assert_numerically_equal( + original, + model, + (np.random.rand(4, 8).astype(np.float32),), + ) + + def test_slice_negative_clamped_start_collapses(self): + """Slice with negative start clamping to 0 (e.g., -10 on dim 8) should collapse.""" + model_text = """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + original = ir.from_onnx_text(model_text) + model = ir.from_onnx_text(model_text) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertIn("Identity", [node.op_type for node in model.graph]) + testing.assert_numerically_equal( + original, + model, + (np.random.rand(4, 8).astype(np.float32),), + ) + + def test_slice_multi_axis_nonzero_start_does_not_collapse(self): + """Slice with matching shapes, unit steps, valid axes, but nonzero start should not collapse.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_single_axis_nonzero_start_does_not_collapse(self): + """Slice with matching shapes, unit steps, valid axis, but nonzero start should not collapse.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_multi_axis_out_of_bounds_axis_does_not_collapse(self): + """Multi-axis slice containing an out-of-bounds axis should not collapse.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0)