A few fixes relating to constant propagation

onnxscript/ir/_core.py

@@ -1612,6 +1612,35 @@
             return
         raise TypeError(f"Expected value to be a Shape or None, got '{type(value)}'")
 
+    def _compute_const_value(self) -> None:
+        """Evaluates Constant nodes and sets the const_value attribute of the Value."""
+        node = self.producer()
+        if node is None:
+            return
+        if node.op_type != "Constant" or node.domain not in {"", "ai.onnx"}:
+            return
+        attr_name, attr_value = next(iter(node.attributes.items()))
+        if attr_value is None or not isinstance(attr_value, Attr):
+            return
+
+        const_value: _protocols.TensorProtocol
+        if attr_name in {"value_float", "value_floats"}:
+            const_value = Tensor(np.array(attr_value.value, dtype=np.float32), name=self.name)
+        elif attr_name in {"value_int", "value_ints"}:
+            const_value = Tensor(np.array(attr_value.value, dtype=np.int64), name=self.name)
+        elif attr_name in {"value_string", "value_strings"}:
+            const_value = StringTensor(
+                np.array(attr_value.value, dtype=np.bytes_), name=self.name
+            )
+        elif attr_name == "value":
+            const_value = typing.cast(_protocols.TensorProtocol, attr_value.value)
+        else:
+            return
+
+        self.const_value = const_value
+        self.shape = const_value.shape  # type: ignore
+        self.dtype = const_value.dtype
+
     @property
     def const_value(
         self,
@@ -1621,6 +1650,10 @@
         The value can be backed by different raw data types, such as numpy arrays.
         The only guarantee is that it conforms TensorProtocol.
         """
+        if self._const_value != None:
+            return self._const_value
+        # On-demand computation of constant value. Currently limited to outputs of Constant nodes.
+        self._compute_const_value()
         return self._const_value
 
     @const_value.setter

onnxscript/optimizer/_constant_folding.py

@@ -32,6 +32,10 @@
     )
 
 
+def is_onnx_op(node: ir.Node, op_type: str) -> bool:
+    return node.op_type == op_type and utils.is_onnx_domain(node.domain)
+
+
 def is_constant_op(node: ir.Node) -> bool:
     return node.op_type in {"Constant", "ConstantOfShape"} and utils.is_onnx_domain(
         node.domain
@@ -168,7 +172,11 @@
         return None
     const_value = val.const_value
     if const_value is not None:
-        return const_value.numpy()
+        try:
+            return const_value.numpy()
+        except FileNotFoundError:
+            # External data is not available.
+            return None
     return None
 
 
@@ -604,6 +612,12 @@
         for i, value in enumerate(node.inputs):
             sym_value = self._state.get_sym_value(value)
             if isinstance(sym_value, ir.Value):
+                logger.debug(
+                    "Node [%s]: Replacing input %s with %s",
+                    node.name,
+                    value.name,
+                    sym_value.name,
+                )
                 node.replace_input_with(i, sym_value)
                 # TODO(rama): consider merging type/other info from both values
 
@@ -626,7 +640,11 @@
                     output = [output]
                 return Replacement(output, context.nodes)
 
-        if is_control_flow_op(node) or is_non_deterministic_op(node):
+        if (
+            is_control_flow_op(node)
+            or is_non_deterministic_op(node)
+            or is_onnx_op(node, "Constant")
+        ):
             return None
 
         input_values = [_get_numpy_value(x) for x in node.inputs]
@@ -648,7 +666,7 @@
             return None
         if len(node.outputs) == 1 and not isinstance(outputs, (tuple, list)):
             replacement = self.new_constant(node.outputs[0], outputs)
-            if is_constant_op(node) or replacement is None:
+            if is_onnx_op(node, "ConstantOfShape") or replacement is None:
                 return None
             return Replacement(replacement.outputs, [replacement])
         else:

onnxscript/rewriter/broadcast_to_matmul.py

@@ -5,7 +5,7 @@
 import logging
 
 from onnxscript import ir
-from onnxscript.rewriter import _ir_utils, pattern
+from onnxscript.rewriter import pattern
 
 logger = logging.getLogger(__name__)
 
@@ -30,8 +30,6 @@ def check_if_not_need_reshape(
 
     input_a_shape = input_a.shape
     input_b_shape = input_b.shape
-    # TODO: Get a helper func to get const_value
-    _ir_utils.propagate_const_value(shape_c)
     shape_c_tensor = shape_c.const_value
     if shape_c_tensor is None:
         logger.info("The value 'shape_c' is not statically known.")

onnxscript/rewriter/onnxruntime/instance_to_group_normalization.py

@@ -7,7 +7,7 @@
 import numpy as np
 import onnx
 
-from onnxscript.rewriter import _ir_utils, pattern
+from onnxscript.rewriter import pattern
 
 torch_module_op = pattern.torch_module_op
 
@@ -42,14 +42,12 @@ def check_if_simulated_instance_norm_is_used(
     Returns:
         bool: True if the simulated instance normalization is used, False otherwise.
     """
-    weight_for_norm_prop = _ir_utils.propagate_const_value(weight_for_norm)
-    weight_for_norm_const_value = weight_for_norm_prop.const_value
+    weight_for_norm_const_value = weight_for_norm.const_value
     if weight_for_norm_const_value is None:
         return False
     weight_for_norm = weight_for_norm_const_value.numpy()
 
-    bias_for_norm_prop = _ir_utils.propagate_const_value(bias_for_norm)
-    bias_for_norm_const_value = bias_for_norm_prop.const_value
+    bias_for_norm_const_value = bias_for_norm.const_value
     if bias_for_norm_const_value is None:
         return False
     bias_for_norm = bias_for_norm_const_value.numpy()
@@ -76,7 +74,6 @@ def check_if_simulated_instance_norm_is_used(
     if not all(dim == 1 for dim in bias_full_shape[1:]):
         return False
 
-    adjusted_input_shape = _ir_utils.propagate_const_value(adjusted_input_shape)
     adjusted_input_shape_const_value = adjusted_input_shape.const_value
 
     g = weight_for_norm.shape[0]
@@ -87,7 +84,6 @@ def check_if_simulated_instance_norm_is_used(
         return False
 
     # NOTE: Restrict the rule to only support constant shape
-    original_input_shape = _ir_utils.propagate_const_value(original_input_shape)
     original_input_shape_const_value = original_input_shape.const_value
     if (
         original_input_shape_const_value is None

onnxscript/rewriter/onnxruntime/transformers/layernorm.py

@@ -5,8 +5,9 @@
 import logging
 
 import onnxscript
+import onnxscript.ir.convenience
 from onnxscript import ir
-from onnxscript.rewriter import _ir_utils, function_rule
+from onnxscript.rewriter import function_rule
 
 logger = logging.getLogger(__name__)
 
@@ -23,7 +24,7 @@ def _fusion(self, function: ir.Function) -> ir.Function:
         if aten_add_node is None:
             raise function_rule.FunctionRewriteError("Could not find Add node")
 
-        eps_ir_value = _ir_utils.propagate_const_value(aten_add_node.inputs[1])
+        eps_ir_value = aten_add_node.inputs[1]
         eps_const_value = eps_ir_value.const_value
         if eps_const_value is None:
             raise function_rule.FunctionRewriteError("Could not find eps")

onnxscript/rewriter/onnxruntime/transformers/multihead_attention.py

@@ -56,8 +56,9 @@
 from onnx import helper as onnx_helper
 
 import onnxscript
+import onnxscript.ir.convenience
 from onnxscript import ir
-from onnxscript.rewriter import _ir_utils, function_rule
+from onnxscript.rewriter import function_rule
 
 logger = logging.getLogger(__name__)
 
@@ -110,7 +111,7 @@ def infer_attn_size_config(self, function: ir.Function) -> AttnSizeConfig:
                 assert (
                     constant_node.op_type == "Constant"
                 ), "Expected the second input to Reshape to be a Constant node."
-                value = _ir_utils.propagate_const_value(reshape_node.inputs[1])
+                value = reshape_node.inputs[1]
                 constant_value = value.const_value
                 if constant_value is None:
                     raise function_rule.FunctionRewriteError(

onnxscript/rewriter/pattern.py

@@ -23,7 +23,6 @@
 
 from onnxscript import ir
 from onnxscript.ir import _convenience, _tape
-from onnxscript.rewriter import _ir_utils
 
 T = TypeVar("T")
 
@@ -618,7 +617,6 @@ def value(self) -> int | float:
         return self._value
 
     def matches(self, value: ir.Value, match: MatchResult) -> MatchResult:
-        value = _ir_utils.propagate_const_value(value)
         constant_value = value.const_value
         if constant_value is None:
             return match.fail(f"Value is not a constant, expecting {self.value}.")
@@ -915,14 +913,16 @@ def _match_constant(self, pattern_constant: Constant, value: ir.Value) -> bool:
         if subgraph replacement happens. But subsequent DCE will remove the constant
         node if it is not used elsewhere.
         """
-        value = _ir_utils.propagate_const_value(value)
         constant_value = value.const_value
         if constant_value is None:
             return self.fail(
                 f"Value {value.name} is not a constant, expecting {pattern_constant.value}.",
             )
 
-        constant_value_numpy = constant_value.numpy()
+        try:
+            constant_value_numpy = constant_value.numpy()
+        except FileNotFoundError:
+            return self.fail(f"Constant value of {value.name} not available.")
         # TODO (rama): allow users to specify shape requirement, if desired.
         if constant_value_numpy.size != 1:
             return self.fail(

onnxscript/rewriter/pattern_test.py

@@ -10,7 +10,7 @@
 
 from onnxscript import FLOAT, ir, script
 from onnxscript import opset17 as op
-from onnxscript.rewriter import _ir_utils, cast_constant_of_shape, pattern
+from onnxscript.rewriter import cast_constant_of_shape, pattern
 
 logger = logging.getLogger(__name__)
 
@@ -259,7 +259,6 @@ def identity(op, x, newshape):
 
         def check_for_redundant_reshape(context, x, newshape):
             oldshape = x.shape
-            newshape = _ir_utils.propagate_const_value(newshape)
             newshape_const_value = newshape.const_value
             if newshape_const_value is None:
                 return False