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

Start in onnxscript.optimizer._constant_folding, especially ReferenceEvaluator.get_evaluator, and compare its operator lookup with the optimized operator table described in the issue. Run the supplied Conv reproduction first; done means constant folding resolves Conv through the optimized implementation and the folded graph remains correct.

_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.