FoldConstantsPass: Why are always_fold_ops restricted to unique consumers?

[iksnagreb]

I do not fully understand the following: Why is always_fold_ops, which I understand as simply overriding the size limit, restricted to operators which are the single unique consumer of an input? Maybe I have missed this, but I do not see a similar restriction for the usual path where the input is below the size limit? Is it to prevent the model size from increasing by implicitly duplicating large tensors via this override?

onnxscript/onnxscript/optimizer/_constant_folding.py

Lines 1037 to 1045 in 413f3df

	if (node.domain, node.op_type) in self.always_fold_ops and all(
	len(input.consumers()) == 1 for input in node.inputs if input is not None
	):
	# If the op is in always_fold_ops and all inputs are used only by this node,
	# we can still fold it even if the input size exceeds the limit.
	logger.debug(
	"Folding large constant for node %s because it is in the always_fold_ops list",
	node,
	)

[justinchuby]

Is it to prevent the model size from increasing by implicitly duplicating large tensors via this override?

I think this is the intention. Suggestions welcome!

[iksnagreb]

Sorry for the delayed answer... I think there is something similar going on here, which took me quite some time to figure out why I cannot get it fold the ConstantOfShape...

onnxscript/onnxscript/optimizer/_constant_folding.py

Lines 1075 to 1076 in 413f3df

	if _is_onnx_op(node, "ConstantOfShape") or replacement is None:
	return None

So at least a debug message would be nice to be aware of this behavior (I wonder if there are even more like this?) Currently you even get something like "DEBUG:onnxscript.optimizer._constant_folding:New constant for value ConstantOfShape_0_output_0 dtype: float32 shape: (1, 1, 15)" when setting the debug log-level, but it then silently decides to just not apply this new constant...

Options to increase the flexibility or to align this closer to user intent could be another option to the FoldConstantsPass besides the always_fold_ops to explicitly allow constant folding with potential model size increase, e.g., allow_size_increase or even more drastic allow_bloat, these could default to false to keep the current behavior. The other option (potentially a breaking change) would be interpret always_fold_ops as ALWAYS fold these operators, even if it increases the size - if someone puts an operator there explicitly, it is probably the intention to really always fold it?

[justinchuby]

Thanks for catching this - for discussion:

I think ConstantOfShape is treated as a constant op (folded). A good behavior may be moving it above the create constant function call so a misleading debug message doesn't get emitted. Would that make sense? Alternatively, when should ConstantOfShape be replaced by constants? Should we create an option for that?

I also like the allow_bloat option.

[gramalingam]

Thanks for bringing up the issue.

(a) Do we need a special behavior for "ConstantOfShape"? I guess, we definitely want a size heuristic (just like any other op): if folding this will increase model size significantly, we don't want to do it. So, the question is: are the general size heuristic options insufficient for "ConstantOfShape"?

I guess the case to consider is when ConstantOfShape produces a tensor of size < output_limit.

(b) I wonder if we are complicating it too much: the documentation of the intended behavior seems complicated as it needs to consider the interaction between the different options to control inlining, with input size limit, output size limit, the bloat option, and always-inline list, etc.

Is there a better alternative? Eg., should we instead allow a dictionary mapping op names to a boolean-function that will return yes or no for the inlining decision?

[justinchuby]

Should the function then be (Node) -> bool, in case input/output shapes etc. needs to be considered? What would be an appropriate amount of information?

[gramalingam]

Yes, Node should be an input ... I guess users would be able to compute sizes of inputs, number of consumers, all from the node ... so, I guess that should be enough to at least encode all of the heuristics we are currently using?