Constant folding uses onnx's plain reference ops, never onnx.reference.ops_optimized

[1duo]

_constant_folding.ReferenceEvaluator.get_evaluator resolves every op through
onnx.reference.ops.load_op(), which only ever returns the plain reference
implementations. onnx's own ReferenceEvaluator defaults to optimized=True and seeds its
op table from onnx.reference.ops_optimized.optimized_operators, so constant folding never
sees the optimized ops.

optimized_operators is just Conv today, but for Conv the difference is im2col+GEMM
versus a Python loop nest, which makes folding a constant Conv orders of magnitude slower
than it needs to be.

Reproduce

import time

import numpy as np
import onnx.reference.ops_optimized
import onnxscript.optimizer
from onnxscript import FLOAT, opset21, script
from onnxscript.optimizer import _constant_folding

X, W = np.ones((1, 2, 192, 192), np.float32), np.ones((32, 2, 1, 1), np.float32)
ATTRS = dict(
    auto_pad="NOTSET", dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0] * 4, strides=[1, 1]
)


@script()
def model(x: FLOAT[1, 32, 192, 192]) -> FLOAT[1, 32, 192, 192]:
    w = opset21.Constant(value=W)
    c = opset21.Constant(value=X)
    return opset21.Add(opset21.Conv(c, w), x)


def timed(label, fn):
    t = time.perf_counter()
    out = fn()
    print(f"{label:<28}{time.perf_counter() - t:7.3f}s")
    return out


picked = _constant_folding.ReferenceEvaluator().get_evaluator("", "Conv", 21)
print("onnxscript resolves Conv to:", picked.__self__.__module__)
print("onnx resolves Conv to:      ", onnx.reference.ops_optimized.Conv.eval.__self__.__module__)

timed("plain Conv.eval", lambda: picked(X, W, **ATTRS))
timed("optimized Conv.eval", lambda: onnx.reference.ops_optimized.Conv.eval(X, W, **ATTRS))
folded = timed(
    "fold_constants",
    lambda: onnxscript.optimizer.fold_constants(
        model.to_model_proto(), input_size_limit=1 << 30, output_size_limit=1 << 30
    ),
)
print("folded graph:", [n.op_type for n in folded.model.graph])

Output on onnxscript 0.7.1 / onnx 1.22.0:

onnxscript resolves Conv to: onnx.reference.ops.op_conv
onnx resolves Conv to:       onnx.reference.ops_optimized.op_conv_optimized
plain Conv.eval               3.314s
optimized Conv.eval           0.005s
fold_constants                3.422s
folded graph: ['Constant', 'Constant', 'Add']

660x on the op itself, and fold_constants spends essentially all of its time there.

Potential fix

Prefer optimized_operators in get_evaluator, matching what ReferenceEvaluator does for
its own new_ops:

_OPTIMIZED_OPS = {cls.__name__: cls for cls in onnx.reference.ops_optimized.optimized_operators}

def get_evaluator(self, domain: str, op: str, version: int) -> Callable | None:
    if not domain and op in _OPTIMIZED_OPS:
        return _OPTIMIZED_OPS[op].eval
    ...

Happy to send a PR if this looks right.

[justinchuby]

Thanks - I think we should just re-implement conv with better vectorized numpy and remove the ops_optimized module. Would that work?

[promiseeuler]

I tested that direction against ONNX 1.22.0 using the reporter's exact X=(1, 2, 192, 192), W=(32, 2, 1, 1) case.

Calling the two underlying implementations directly produced identical values (np.allclose == True, max absolute error 0.0). On this Apple Silicon machine, the plain loop took 4.98 s and _conv_implementation_im2col took 5.7 ms. The optimized Conv._run is otherwise the same wrapper and casts the result back to X.dtype, so moving the vectorized helper into onnx.reference.ops.op_conv.Conv would solve this for ONNXScript without changing ReferenceEvaluator.get_evaluator(): its existing load_op("", "Conv", version) call would then receive the fast implementation.

The caveat is release ownership/timing: that change and removal/deprecation of onnx.reference.ops_optimized belong in the ONNX repository, and ONNXScript would remain slow for supported ONNX versions before that release. If ONNXScript needs the improvement immediately, preferring optimized_operators here is still the local compatibility bridge.

For the upstream replacement, I would retain parity tests across 1D/2D/3D convolution, groups, bias/no-bias, asymmetric pads, SAME_*/VALID, strides, dilations, and float dtypes. One subtlety from the direct helpers: the loop returned float64 and im2col returned float32 in this benchmark, although both public _run wrappers normalize back to X.dtype.

[MohammedAlkindi]

Independent x86-64 confirmation, since the measurements so far are Apple Silicon.

Windows 11, Intel, Python 3.13.13, onnx 1.22.0, onnxscript 0.7.1, numpy 2.5.2. get_evaluator("", "Conv", 21) resolves to onnx.reference.ops.op_conv, while onnx's optimized Conv is ops_optimized.op_conv_optimized. On the reporter's exact case, plain .eval() takes 5.19s against 0.0106s optimized, a 490x difference, with allclose True and max absolute error 0.0. Same order as the 4.98s/5.7ms reported on Apple Silicon, so this is not architecture-specific.

Promise Emmanuel Oluwadare (@promiseeuler)'s dtype point reproduces here, and it matters for the proposed move: the direct helpers differ (_conv_implementation returns float64, _conv_implementation_im2col returns float32), while both public .eval() wrappers normalize to float32. So relocating the im2col helper into op_conv would change the internal dtype for anything calling the helper directly, invisibly through _run.