Constant-folding behavior inside functions

[iksnagreb]

With the recent change to the constant-folding (see #2650) I started to encounter problems with constant folding within functions: The model checker complains about topological sorting and onnxruntime complains about "Input 0 expected to have type but instead is null" afterwards, but trying to narrow this down it seems the root cause is an initializer created by constant folding is missing from the graph. The following model should be a minimal example to reproduce this:

<
   ir_version: 9,
   opset_import: ["this" : 1, "" : 19]
>
model (float x) => (float return_val) {
   [n0] return_val = this.function (x)
}
<
  domain: "this",
  opset_import: ["" : 19]
>
function (x) => (return_val)
{
   [n0] tmp = Constant <value_int: int = 1> ()
   [n1] tmp_0 = Cast <to: int = 1> (tmp)
   [n2] return_val = Sub (tmp_0, x)
}

Up until including v0.5.4 this folds (correctly?) to the following, where a single constant node replaces the cast operator:

<
   ir_version: 9,
   opset_import: ["this" : 1, "" : 19]
>
model (float x) => (float return_val) 
   <float "this::function/tmp_0">
{
   [n0] return_val = this.function (x)
}
<
  domain: "this",
  opset_import: ["" : 19]
>
function (x) => (return_val)
{
   [node_Constant_0] tmp_0 = Constant <value: tensor = float tmp_0 {1}> ()
   [n2] return_val = Sub (tmp_0, x)
}

Since v0.5.5, however, the resulting model looks as follows, where the temporary has no producer and also no initializer (which I would expect according to the recent changes):

<
   ir_version: 9,
   opset_import: ["this" : 1, "" : 19]
>
model (float x) => (float return_val) {
   [n0] return_val = this.function (x)
}
<
  domain: "this",
  opset_import: ["" : 19]
>
function (x) => (return_val)
{
   [n2] return_val = Sub (tmp_0, x)
}

The model checker raises ValidationError: Nodes in a function must be topologically sorted, however input 'tmp_0' of node: Name: n2 OpType: Sub is neither output of any previous nodes nor input of the function. The following reproduces the two outputs, depending on the onnxscript version:

import onnx
import onnx_ir as ir

import onnxscript.optimizer as optimizer
from onnxscript.optimizer import _constant_folding

model = """
<
   ir_version: 9,
   opset_import: ["this" : 1, "" : 19]
>
model (float x) => (float return_val) {
   [n0] return_val = this.function (x)
}
<
  domain: "this",
  opset_import: ["" : 19]
>
function (x) => (return_val)
{
   [n0] tmp = Constant <value_int: int = 1> ()
   [n1] tmp_0 = Cast <to: int = 1> (tmp)
   [n2] return_val = Sub (tmp_0, x)
}
"""
model = ir.from_onnx_text(model)
_constant_folding.fold_constants(model)
optimizer.remove_unused_nodes(model)
print(ir.to_onnx_text(model))
onnx.checker.check_model(ir.serde.serialize_model(model))

Is this a bug, or, as #2650 was labeled as a breaking change, expected behavior and I need to adjust my models somehow?

[justinchuby]

Looks like a bug because in functions we should keep creating constant nodes. @titaiwangms

[iksnagreb]

Thanks!