ONNXScript Constant Folding: A Bug Or A Feature?

Hey guys! Let's talk about something a bit technical, but super important if you're working with ONNX and ONNXScript: constant folding within functions. Recently, there have been some changes, and as a result, some of us have encountered issues. This article dives deep into the problem, exploring the behavior, potential bugs, and how it impacts your models.

The Problem: Missing Initializers and Topological Sorting Errors

So, what's the deal? The core issue revolves around constant folding within functions in ONNXScript. Specifically, after a recent change (referenced as #2650), models are experiencing problems with topological sorting and errors like "Input 0 expected to have type but instead is null" from onnxruntime. The main culprit? It seems that an initializer created during constant folding is sometimes missing from the graph, leading to these validation errors. Let's break it down to understand what’s happening in detail. We'll examine a minimal example that replicates the problem.

Imagine you have a function that looks like this (in a simplified ONNXScript representation):

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)
}

In this function, you're creating a constant integer value (1), casting it to a float, and then subtracting the input 'x' from it. The goal of constant folding is to optimize this by replacing the 'Cast' operation with a pre-calculated constant value. This is a type of optimization that simplifies the model and can improve performance. However, with the changes, something is going wrong. The model checker is complaining about nodes not being topologically sorted, meaning the order of operations is incorrect, and onnxruntime is getting confused because it expects a value for 'tmp_0' that isn't there.

The implications of these errors can be significant, potentially causing models to fail during validation or runtime. This is a critical issue that developers using ONNXScript need to be aware of, as it can affect the correctness and performance of their models. We must understand how these changes impact the constant folding process. Let's explore the before and after scenarios to get a clearer picture of this problem.

Before v0.5.5: Expected Behavior

Before version v0.5.5 of ONNXScript, the constant folding process appeared to work as expected. The code would correctly fold, which means it would simplify the operation within the function. After the folding, the graph would look like this (again, simplified):

<float "this::function/tmp_0">
{
   [node_Constant_0] tmp_0 = Constant <value: tensor = float tmp_0 {1}> ()
   [n2] return_val = Sub (tmp_0, x)
}

Essentially, the 'Cast' operator is replaced by a pre-calculated constant node. The output of the 'Cast' operator is then replaced by the constant value directly. This optimization eliminates an unnecessary operation, making the model more efficient. This is the expected outcome of constant folding – simplifying the graph and reducing the number of computations required.

In the pre-v0.5.5 scenario, the initializer for 'tmp_0' (the constant value) would be present in the graph, ensuring that onnxruntime could find it and perform the subtraction correctly. There were no topological sorting issues, and everything worked as intended. The model was optimized, and there were no unexpected errors. This behavior allowed for a smooth and efficient execution of the ONNX models. Now, let's explore what happened when ONNXScript was updated to v0.5.5 to get a better understanding of the problem.

After v0.5.5: The Bug Appears?

However, things change with v0.5.5. This is where the issues arise. The model generated after constant folding now looks like this:

function (x) => (return_val)
{
   [n2] return_val = Sub (tmp_0, x)
}

The crucial difference here is the absence of the initializer for 'tmp_0'. The constant value is not present in the graph. The subtraction operation now directly refers to 'tmp_0' without it having a producer or a definition. This leads to the validation errors.

The model checker raises a 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." This error indicates that the 'Sub' node is trying to use a value ('tmp_0') that hasn't been computed or defined earlier in the graph. Since there's no initializer and no preceding node that produces 'tmp_0', the topological order is violated.

This behavior suggests that the constant folding process is not correctly handling constants within functions in v0.5.5, leading to missing initializers and validation failures. This is a significant regression, impacting models that rely on constant folding for optimization. So, the big question is: is this a bug, or is it intended behavior due to the changes in #2650? Let's consider that question next.

Bug or Feature? The Dilemma

Is this a bug, or is it the new normal? The question is critical. If it's a bug, it needs to be fixed. If it's intended behavior, developers need to adjust their models to accommodate the changes. The original poster points out that #2650 was labeled as a breaking change. This suggests that the behavior might be intentional. However, the resulting behavior breaks existing models and introduces significant compatibility issues. It is important to know whether this is going to be the new behavior or if it’s a bug that needs to be addressed.

If it’s a feature, we need to understand how to adapt our models. Do we need to modify our ONNXScript code to ensure that constants are handled correctly? Are there new best practices for writing models that use constant folding? This is something the community needs to discuss, so we can ensure we stay on the right track.

The community must come together and work with the ONNXScript maintainers to clarify the situation. If it is a bug, the developers need to pinpoint the root cause and implement a fix, ensuring that constant folding functions correctly. If this is a feature, we need to adapt to a new paradigm. It is important to share information, test models, and provide feedback to help shape the future of ONNXScript. The answer to this question shapes the future of the ONNX models.

Reproducing the Issue: Code Example

To better understand and verify this issue, here is the Python code provided in the original post. This code demonstrates the problem. When you run this code using different versions of ONNXScript, you'll see the difference. This difference in behavior is due to the constant folding issue. This allows for easier testing of the issue.

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))

This script takes the ONNX model, performs constant folding, removes any unused nodes, and then attempts to validate the resulting model. The key lies in the _constant_folding.fold_constants(model) line. This is where the problematic behavior is triggered. By running this script with various ONNXScript versions, you can witness the discrepancy. This helps to showcase the problem in action.

Conclusion: A Call to Action

In conclusion, the constant folding behavior within functions in ONNXScript v0.5.5 and later introduces a potential issue. Missing initializers and topological sorting errors result in significant challenges for developers. Whether it's a bug or a feature is yet to be determined, but the community must come together to determine a plan. Understanding the impact of this change on existing models and the need for potential adjustments is crucial for the ongoing development and use of ONNXScript.

If you're using ONNXScript, I encourage you to test your models and keep a close eye on this issue. Share your findings, and engage in discussions to help clarify the situation. By working together, we can ensure that ONNXScript remains a reliable and powerful tool for the future. The future of the ONNX models lies in these collaborative efforts. Together, we can make the most out of ONNXScript.