diff --git a/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py b/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py index 8c63d3e777..51c8231143 100644 --- a/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py +++ b/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py @@ -18,13 +18,30 @@ def cast_constant_of_shape(op, shape, scalar, dtype): def fused_cast_constant_of_shape(op, shape: ir.Value, scalar: ir.Attr, dtype: ir.Attr, **_): # Cast scalar (a TensorProto attribute) to the specified dtype - scalar_value = scalar.value.numpy().item() + scalar_value = scalar.as_tensor().numpy().item() cast_value = ir.tensor([scalar_value], dtype=ir.DataType(dtype.as_int())) return op.ConstantOfShape(shape, value=cast_value) +def check_cast_constant_of_shape(context, shape, scalar, dtype, **_) -> bool: + del context, shape + if dtype.is_ref() or scalar.is_ref(): + return False + if dtype.type != ir.AttributeType.INT: + return False + to_dtype = dtype.as_int() + if to_dtype in (ir.DataType.UNDEFINED, ir.DataType.STRING): + return False + if scalar.type != ir.AttributeType.TENSOR or scalar.value is None: + return False + scalar_tensor = scalar.as_tensor() + if scalar_tensor.size != 1 or scalar_tensor.dtype == ir.DataType.STRING: + return False + return True + + def cast_constant_of_shape_without_value(op, shape, dtype): - constant = op.ConstantOfShape(shape) + constant = op.ConstantOfShape(shape, _allow_other_attributes=False) return op.Cast(constant, to=dtype) @@ -33,10 +50,28 @@ def fused_cast_constant_of_shape_without_value(op, shape, dtype, **_): return op.ConstantOfShape(shape, value=zero) -cast_constant_of_shape_rule = RewriteRule(cast_constant_of_shape, fused_cast_constant_of_shape) +def check_cast_constant_of_shape_without_value(context, shape, dtype, **_) -> bool: + del context, shape + if dtype.is_ref(): + return False + if dtype.type != ir.AttributeType.INT: + return False + to_dtype = dtype.as_int() + if to_dtype in (ir.DataType.UNDEFINED, ir.DataType.STRING): + return False + return True + + +cast_constant_of_shape_rule = RewriteRule( + cast_constant_of_shape, + fused_cast_constant_of_shape, + check_cast_constant_of_shape, +) cast_constant_of_shape_without_value_rule = RewriteRule( - cast_constant_of_shape_without_value, fused_cast_constant_of_shape_without_value + cast_constant_of_shape_without_value, + fused_cast_constant_of_shape_without_value, + check_cast_constant_of_shape_without_value, ) rules = RewriteRuleSet( diff --git a/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py b/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py index 9001f95806..a6e4b82a21 100644 --- a/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py +++ b/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py @@ -1,11 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +from __future__ import annotations + import unittest +import numpy as np import onnx.checker import onnx.parser import onnx_ir as ir +from onnxscript.rewriter import testing from onnxscript.rewriter.rules.common import _cast_constant_of_shape @@ -49,6 +53,181 @@ def test_cast_after_constant_of_shape_without_value_is_fused(self): output_model_proto = ir.serde.serialize_model(model) onnx.checker.check_model(output_model_proto, True) + def test_cast_to_string_after_constant_of_shape_is_not_fused(self): + """Cast to string should not be fused into ConstantOfShape as ConstantOfShape does not support STRING output.""" + model_text = """ + + agraph (int64[2] input_x) => (string[1, 4] output) + { + constant = ConstantOfShape (input_x) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + # Verify the model serializes cleanly without SerdeError + output_proto = ir.serde.serialize_model(model) + self.assertIsNotNone(output_proto) + + def test_cast_to_string_after_constant_of_shape_without_value_is_not_fused(self): + """Cast to string without value should not be fused into ConstantOfShape as ConstantOfShape does not support STRING output.""" + model_text = """ + + agraph (int64[2] input_x) => (string[1, 4] output) + { + constant = ConstantOfShape (input_x) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + # Verify the model serializes cleanly without SerdeError + output_proto = ir.serde.serialize_model(model) + self.assertIsNotNone(output_proto) + + def test_cast_after_constant_of_shape_with_ref_attr_is_not_fused(self): + """Cast after ConstantOfShape with a reference value attribute should not be fused or replaced with zero.""" + model_text = """ + + main_graph (int64[2] shape) => (float[1, 4] output) { + output = pkg.my_func (shape) + } + + my_func (shape) => (output) { + constant = ConstantOfShape (shape) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + function = next(iter(model.functions.values())) + self.assertEqual([node.op_type for node in function], ["ConstantOfShape", "Cast"]) + constant_node = function[0] + self.assertIn("value", constant_node.attributes) + self.assertTrue(constant_node.attributes["value"].is_ref()) + self.assertEqual(constant_node.attributes["value"].ref_attr_name, "fill") + + def test_cast_after_constant_of_shape_with_non_scalar_value_is_not_fused(self): + """Cast after ConstantOfShape with a non-scalar value attribute should not be fused.""" + model_text = """ + + agraph (int64[2] input_x) => (float[1, 4] output) + { + constant = ConstantOfShape (input_x) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + self.assertEqual([node.op_type for node in model.graph], ["ConstantOfShape", "Cast"]) + + def test_cast_after_constant_of_shape_without_value_numerical_correctness(self): + """Verify numerical correctness of legitimate ConstantOfShape without value cast fusion.""" + model_text = """ + + agraph (int64[2] input_x) => (float[1, 4] output) + { + constant = ConstantOfShape (input_x) + output = Cast (constant) + } + """ + original = ir.from_onnx_text(model_text) + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 1) + testing.assert_numerically_equal( + original, + model, + (np.array([1, 4], dtype=np.int64),), + ) + + def test_cast_with_ref_attr_to_after_constant_of_shape_is_not_fused(self): + """Cast with reference target dtype attribute should not be fused.""" + model_text = """ + + main_graph (int64[2] shape) => (float[1, 4] output) { + output = pkg.my_func (shape) + } + + my_func (shape) => (output) { + constant = ConstantOfShape (shape) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + function = next(iter(model.functions.values())) + self.assertEqual([node.op_type for node in function], ["ConstantOfShape", "Cast"]) + + def test_cast_with_ref_attr_to_after_constant_of_shape_without_value_is_not_fused(self): + """Cast with reference target dtype without value should not be fused.""" + model_text = """ + + main_graph (int64[2] shape) => (float[1, 4] output) { + output = pkg.my_func (shape) + } + + my_func (shape) => (output) { + constant = ConstantOfShape (shape) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + function = next(iter(model.functions.values())) + self.assertEqual([node.op_type for node in function], ["ConstantOfShape", "Cast"]) + + def test_cast_to_undefined_after_constant_of_shape_is_not_fused(self): + """Cast to undefined data type should not be fused into ConstantOfShape.""" + model_text = """ + + agraph (int64[2] input_x) => (float[1, 4] output) + { + constant = ConstantOfShape (input_x) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + self.assertEqual([node.op_type for node in model.graph], ["ConstantOfShape", "Cast"]) + + def test_cast_to_undefined_after_constant_of_shape_without_value_is_not_fused(self): + """Cast to undefined data type without value should not be fused into ConstantOfShape.""" + model_text = """ + + agraph (int64[2] input_x) => (float[1, 4] output) + { + constant = ConstantOfShape (input_x) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + self.assertEqual([node.op_type for node in model.graph], ["ConstantOfShape", "Cast"]) + + def test_cast_after_constant_of_shape_with_string_value_is_not_fused(self): + """Cast after ConstantOfShape with a string value tensor should not be fused.""" + model_text = """ + + agraph (int64[2] input_x) => (float[1, 4] output) + { + constant = ConstantOfShape (input_x) + output = Cast (constant) + } + """ + model = ir.from_onnx_text(model_text) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 0) + self.assertEqual([node.op_type for node in model.graph], ["ConstantOfShape", "Cast"]) + if __name__ == "__main__": unittest.main()