diff --git a/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py b/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py index 09c5db7735..d7621f76ea 100644 --- a/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py +++ b/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py @@ -80,9 +80,15 @@ def pattern(self, op, data, indices, updates): def check(self, context, data, indices, updates, **_): """Check if the ScatterND is redundant due to static indices covering entire tensor.""" - # To validate data can be replaced directly by updates, we need to check the following: - # 1. they have the same shape result = onnxscript.rewriter.MatchResult() + # To validate data can be replaced directly by updates, we need to check the following: + # 1. the scatter-update is an assignment, not a reduction combining data and updates + reduction = context.root.attributes.get_string("reduction", "none") + if reduction != "none": + return result.fail( + f"The 'reduction' is {reduction!r}, so 'data' still contributes to the result." + ) + # 2. they have the same shape if data.shape is None: return result.fail("The value 'data' shape is not statically known.", data) if updates.shape is None: @@ -92,7 +98,7 @@ def check(self, context, data, indices, updates, **_): "The shape of 'data' and 'updates' are different.", [data, updates] ) - # 2. the indices is referring to the whole data, which is from 0 to data.shape[0] + # 3. the indices is referring to the whole data, which is from 0 to data.shape[0] if indices.const_value is None: return result.fail("The value 'indices' is not statically known.", indices) expected_indices = [[i] for i in range(data.shape[0])] diff --git a/onnxscript/rewriter/rules/common/_redundant_scatter_nd_test.py b/onnxscript/rewriter/rules/common/_redundant_scatter_nd_test.py index 96e3bcc80c..cc45066cf7 100644 --- a/onnxscript/rewriter/rules/common/_redundant_scatter_nd_test.py +++ b/onnxscript/rewriter/rules/common/_redundant_scatter_nd_test.py @@ -8,6 +8,7 @@ import onnx.parser import onnx_ir as ir import onnxruntime +import parameterized from onnx_ir.passes.common import CheckerPass, ShapeInferencePass import onnxscript.optimizer @@ -120,6 +121,53 @@ def test_redundant_scatter_nd_static_indices(self): for original_output, optimized_output in zip(original_outputs, optimized_outputs): np.testing.assert_allclose(original_output, optimized_output, rtol=1e-6, atol=1e-6) + @parameterized.parameterized.expand( + [ + ("add",), + ("mul",), + ("max",), + ("min",), + ] + ) + def test_static_indices_with_reduction_is_not_redundant(self, reduction: str): + """A reducing ScatterND combines data with updates, so it is not an Identity.""" + model_proto = onnx.parser.parse_model( + f""" + + agraph (float[8, 16] data, float[8, 16] updates) => (float[8, 16] output) + {{ + output = ScatterND (data, indices, updates) + }} + """ + ) + indices = np.arange(8).reshape(8, 1).astype(np.int64) + model = ir.serde.deserialize_model(model_proto) + indices_value = model.graph[0].inputs[1] + indices_value.const_value = ir.Tensor(name="indices", value=indices) + model.graph.initializers["indices"] = indices_value + original_model_proto = ir.serde.serialize_model(model) + + _redundant_scatter_nd.rules.apply_to_model(model) + optimized_model_proto = ir.serde.serialize_model(model) + + # Test numerical equivalence + inputs = { + "data": np.random.rand(8, 16).astype(np.float32), + "updates": np.random.rand(8, 16).astype(np.float32), + } + session = onnxruntime.InferenceSession( + original_model_proto.SerializeToString(), providers=["CPUExecutionProvider"] + ) + original_outputs = session.run(None, inputs) + optimized_session = onnxruntime.InferenceSession( + optimized_model_proto.SerializeToString(), providers=["CPUExecutionProvider"] + ) + optimized_outputs = optimized_session.run(None, inputs) + + for original_output, optimized_output in zip(original_outputs, optimized_outputs): + np.testing.assert_allclose(original_output, optimized_output, rtol=1e-6, atol=1e-6) + self.assertIn("ScatterND", [node.op_type for node in model.graph]) + if __name__ == "__main__": unittest.main()