Skip to content

Add guard to reject apply_to_subgraphs=True for self-subgraph-handling transforms #240

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
auphelia wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add guard to reject apply_to_subgraphs=True for self-subgraph-handling transforms #240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 48 additions & 103 deletions src/qonnx/core/modelwrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,15 +31,15 @@

import copy
import inspect
from pathlib import Path
import numpy as np
import onnx
import onnx.helper as oh
import onnx.numpy_helper as np_helper
import os
import warnings
import numpy as np
from typing import TYPE_CHECKING, Any, Callable, Sequence, TypeVar, cast
from onnx import GraphProto, ModelProto, NodeProto, TensorProto, ValueInfoProto
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Sequence, TypeVar, cast

if TYPE_CHECKING:
from qonnx.transformation.base import Transformation
Expand Down Expand Up @@ -81,17 +81,13 @@ def __init__(
initializers that are missing theirs.
"""
if isinstance(onnx_model_proto, str):
assert os.path.isfile(onnx_model_proto), (
f"File not found: {onnx_model_proto}"
)
assert os.path.isfile(onnx_model_proto), f"File not found: {onnx_model_proto}"
self._model_proto: ModelProto = onnx.load(onnx_model_proto)
elif isinstance(onnx_model_proto, bytes):
self._model_proto: ModelProto = onnx.load_from_string(onnx_model_proto)
else:
if make_deepcopy:
self._model_proto: ModelProto = copy.deepcopy(
cast(ModelProto, onnx_model_proto)
)
self._model_proto: ModelProto = copy.deepcopy(cast(ModelProto, onnx_model_proto))
else:
self._model_proto: ModelProto = cast(ModelProto, onnx_model_proto)
self.temporary_fix_oldstyle_domain()
Expand Down Expand Up @@ -150,15 +146,14 @@ def save(self, filename: str | Path) -> None:

def analysis(
self,
analysis_fxn: Callable[["ModelWrapper", bool], T]
| Callable[["ModelWrapper"], T],
analysis_fxn: Callable[["ModelWrapper", bool], T] | Callable[["ModelWrapper"], T],
apply_to_subgraphs: bool = False,
) -> T:
"""Runs given anaylsis_fxn on this model and return resulting dict."""
if apply_to_subgraphs:
assert "apply_to_subgraphs" in inspect.signature(analysis_fxn).parameters, (
"analysis_fxn must have 'apply_to_subgraphs' argument when apply_to_subgraphs == True"
)
assert (
"apply_to_subgraphs" in inspect.signature(analysis_fxn).parameters
), "analysis_fxn must have 'apply_to_subgraphs' argument when apply_to_subgraphs == True"
return analysis_fxn(self, apply_to_subgraphs) # type: ignore[call-arg]
else:
return analysis_fxn(self) # type: ignore
Expand Down Expand Up @@ -222,13 +217,23 @@ def transform(
- cleanup : execute cleanup transformations before returning
- apply_to_subgraphs : if True, transformation is applied to all subgraphs of the model
"""
# Some transformations manage their own subgraph traversal (e.g. because they
# need hierarchical context that the generic descent cannot provide). Such
# transformations set handles_subgraphs_internally = True and must not be
# driven via the generic apply_to_subgraphs=True path, which re-enters each
# subgraph as a standalone top-level model and loses that context.
if apply_to_subgraphs and getattr(transformation, "handles_subgraphs_internally", False):
raise ValueError(
f"{type(transformation).__name__} manages subgraph traversal itself; "
"do not pass apply_to_subgraphs=True to transform(). Configure subgraph "
"handling on the transform instead."
)

transformed_model = self
if make_deepcopy:
transformed_model = copy.deepcopy(self)
if self.fix_float64:
(transformed_model, model_was_changed) = DoubleToSingleFloat().apply(
transformed_model
)
(transformed_model, model_was_changed) = DoubleToSingleFloat().apply(transformed_model)

if apply_to_subgraphs and (use_preorder_traversal is False):
transformed_model.transform_subgraphs(
Expand All @@ -241,9 +246,7 @@ def transform(

model_was_changed = True
while model_was_changed:
(transformed_model, model_was_changed) = transformation.apply(
transformed_model
)
(transformed_model, model_was_changed) = transformation.apply(transformed_model)
if cleanup:
transformed_model.cleanup()

Expand All @@ -269,27 +272,19 @@ def cleanup(self) -> ModelWrapper:
GiveUniqueParameterTensors(),
]
for trn in cleanup_transforms:
transformed_model = transformed_model.transform(
trn, cleanup=False, make_deepcopy=False
)
transformed_model = transformed_model.transform(trn, cleanup=False, make_deepcopy=False)
return transformed_model

def make_subgraph_modelwrapper(self, subgraph: GraphProto) -> ModelWrapper:
return ModelWrapper(
util.qonnx_make_model(
subgraph, opset_imports=self._model_proto.opset_import
)
)
return ModelWrapper(util.qonnx_make_model(subgraph, opset_imports=self._model_proto.opset_import))

def get_tensor_datatype(self, tensor_name: str) -> BaseDataType:
"""Returns the QONNX DataType of tensor with given name."""
graph = self._model_proto.graph
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret = util.get_by_name(
ret.quant_parameter_tensor_names, "finn_datatype", "key"
)
ret = util.get_by_name(ret.quant_parameter_tensor_names, "finn_datatype", "key")
if ret is not None:
return DataType[ret.value]
onnx_dtype_to_qonnx_dtype = {
Expand Down Expand Up @@ -323,9 +318,7 @@ def set_tensor_datatype(self, tensor_name: str, datatype: BaseDataType | None) -
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret_dt = util.get_by_name(
ret.quant_parameter_tensor_names, "finn_datatype", "key"
)
ret_dt = util.get_by_name(ret.quant_parameter_tensor_names, "finn_datatype", "key")
if ret_dt is not None:
if datatype is None:
ret_dt.Clear()
Expand Down Expand Up @@ -353,18 +346,14 @@ def get_tensor_valueinfo(self, tensor_name: str) -> ValueInfoProto | None:
vi_names += [(x.name, x) for x in graph.value_info]
try:
vi_t_names = [x[0] for x in vi_names]
assert vi_t_names.count(tensor_name) <= 1, (
"Multiple ValueInfoProto found for " + tensor_name
)
assert vi_t_names.count(tensor_name) <= 1, "Multiple ValueInfoProto found for " + tensor_name
vi_ind = vi_t_names.index(tensor_name)
vi = vi_names[vi_ind][1]
return vi
except ValueError:
return None

def get_tensor_shape(
self, tensor_name: str, fix_missing_init_shape: bool = False
) -> list[int] | None:
def get_tensor_shape(self, tensor_name: str, fix_missing_init_shape: bool = False) -> list[int] | None:
"""Returns the shape of tensor with given name, if it has ValueInfoProto.
If fix_missing_init_shape is specified, it will add a ValueInfoProto for initializers
that are missing theirs."""
Expand All @@ -374,33 +363,25 @@ def get_tensor_shape(
vi_names += [(x.name, x) for x in graph.value_info]
try:
vi_t_names = [x[0] for x in vi_names]
assert vi_t_names.count(tensor_name) <= 1, (
"Multiple ValueInfoProto found for " + tensor_name
)
assert vi_t_names.count(tensor_name) <= 1, "Multiple ValueInfoProto found for " + tensor_name
vi_ind = vi_t_names.index(tensor_name)
vi = vi_names[vi_ind][1]
dims = [x.dim_value for x in vi.type.tensor_type.shape.dim]
return dims
except ValueError:
# no ValueInfo found for tensor, check initializer
# (see https://github.com/onnx/onnx/issues/2874)
tensor_init, tensor_init_dtype = self.get_initializer(
tensor_name, return_dtype=True
) # type: ignore
tensor_init, tensor_init_dtype = self.get_initializer(tensor_name, return_dtype=True) # type: ignore
if tensor_init is None:
# no shape defined for this tensor
return None
else:
if fix_missing_init_shape:
self.set_tensor_shape(
tensor_name, tensor_init.shape, dtype=tensor_init_dtype
)
self.set_tensor_shape(tensor_name, tensor_init.shape, dtype=tensor_init_dtype)
# use list return type to keep it consistent with ValueInfo case
return list(tensor_init.shape)

def set_tensor_shape(
self, tensor_name: str, tensor_shape: Sequence[int], dtype: int | None = None
) -> None:
def set_tensor_shape(self, tensor_name: str, tensor_shape: Sequence[int], dtype: int | None = None) -> None:
"""Assigns shape in ValueInfoProto for tensor with given name. If override_dtype
is None, it will try to preserve the existing datatype, otherwise defaults to
single-precision float."""
Expand Down Expand Up @@ -478,9 +459,7 @@ def rename_tensor(self, old_name: str, new_name: str) -> None:
if initializer is not None:
initializer.name = new_name
# sweep over quantization annotations
quant_annotation = util.get_by_name(
graph.quantization_annotation, old_name, "tensor_name"
)
quant_annotation = util.get_by_name(graph.quantization_annotation, old_name, "tensor_name")
if quant_annotation is not None:
quant_annotation.tensor_name = new_name
# sweep over node i/o
Expand Down Expand Up @@ -564,9 +543,7 @@ def find_consumer(self, tensor_name: str) -> NodeProto | None:
elif len(ret) == 1:
return ret[0]
else:
warnings.warn(
"find_consumer: found multiple consumers, returning first one"
)
warnings.warn("find_consumer: found multiple consumers, returning first one")
return ret[0]

def find_consumers(self, tensor_name: str) -> list[NodeProto]:
Expand Down Expand Up @@ -613,13 +590,9 @@ def is_fork_node(self, node: NodeProto) -> bool:
# if the node output is also wired to a top-level output, it is still
# a fork with only 1 direct successor
if node.output[0] in [x.name for x in self.graph.output]:
is_fork = (
False if direct_successors is None else (len(direct_successors) > 0)
)
is_fork = False if direct_successors is None else (len(direct_successors) > 0)
else:
is_fork = (
False if direct_successors is None else (len(direct_successors) > 1)
)
is_fork = False if direct_successors is None else (len(direct_successors) > 1)
return is_fork

def is_join_node(self, node: NodeProto) -> bool:
Expand All @@ -629,13 +602,9 @@ def is_join_node(self, node: NodeProto) -> bool:
# if the node input is also wired to a top-level input, it is still
# a fork with only 1 direct predecessor
if node.input[0] in [x.name for x in self.graph.input]:
is_join = (
False if direct_predecessors is None else (len(direct_predecessors) > 0)
)
is_join = False if direct_predecessors is None else (len(direct_predecessors) > 0)
else:
is_join = (
False if direct_predecessors is None else (len(direct_predecessors) > 1)
)
is_join = False if direct_predecessors is None else (len(direct_predecessors) > 1)
return is_join

def get_all_tensor_names(self) -> list[str]:
Expand Down Expand Up @@ -690,9 +659,7 @@ def make_empty_exec_context(self) -> dict[str, np.ndarray | None]:
execution_context[""] = None
return execution_context

def check_all_tensor_shapes_specified(
self, fix_missing_init_shape: bool = False
) -> bool:
def check_all_tensor_shapes_specified(self, fix_missing_init_shape: bool = False) -> bool:
"""Checks whether all tensors have a specified shape (ValueInfo).
The ONNX standard allows for intermediate activations to have no
associated ValueInfo, but QONNX expects this.
Expand All @@ -707,19 +674,9 @@ def check_all_tensor_shapes_specified(
for i in n.input:
# skip tensor names with empty string (indicates defaults)
if i != "":
ret = (
self.get_tensor_shape(
i, fix_missing_init_shape=fix_missing_init_shape
)
is not None
) and ret
ret = (self.get_tensor_shape(i, fix_missing_init_shape=fix_missing_init_shape) is not None) and ret
for o in n.output:
ret = (
self.get_tensor_shape(
o, fix_missing_init_shape=fix_missing_init_shape
)
is not None
) and ret
ret = (self.get_tensor_shape(o, fix_missing_init_shape=fix_missing_init_shape) is not None) and ret
return ret

def get_tensor_fanout(self, tensor_name: str) -> int:
Expand Down Expand Up @@ -801,9 +758,7 @@ def get_tensor_layout(self, tensor_name: str) -> list[str] | None:
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret = util.get_by_name(
ret.quant_parameter_tensor_names, "tensor_layout", "key"
)
ret = util.get_by_name(ret.quant_parameter_tensor_names, "tensor_layout", "key")
if ret is not None:
return eval(ret.value)
return None
Expand All @@ -816,9 +771,7 @@ def set_tensor_layout(self, tensor_name: str, data_layout: list[str]) -> None:
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret_tl = util.get_by_name(
ret.quant_parameter_tensor_names, "tensor_layout", "key"
)
ret_tl = util.get_by_name(ret.quant_parameter_tensor_names, "tensor_layout", "key")
if ret_tl is not None:
ret_tl.value = str(data_layout)
else:
Expand All @@ -841,24 +794,18 @@ def get_tensor_sparsity(self, tensor_name: str) -> dict[str, Any] | None:
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret = util.get_by_name(
ret.quant_parameter_tensor_names, "tensor_sparsity", "key"
)
ret = util.get_by_name(ret.quant_parameter_tensor_names, "tensor_sparsity", "key")
if ret is not None:
return eval(ret.value)
return None

def set_tensor_sparsity(
self, tensor_name: str, sparsity_dict: dict[str, Any]
) -> None:
def set_tensor_sparsity(self, tensor_name: str, sparsity_dict: dict[str, Any]) -> None:
"""Sets the sparsity annotation of a tensor with given name."""
graph = self._model_proto.graph
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret_ts = util.get_by_name(
ret.quant_parameter_tensor_names, "tensor_sparsity", "key"
)
ret_ts = util.get_by_name(ret.quant_parameter_tensor_names, "tensor_sparsity", "key")
if ret_ts is not None:
ret_ts.value = str(sparsity_dict)
else:
Expand All @@ -879,9 +826,7 @@ def get_opset_imports(self):
"""Returns a list of imported opsets as a {domain, version} dictionary."""
return {opset.domain: opset.version for opset in self._model_proto.opset_import}

def get_customop_wrapper(
self, node, fallback_customop_version=util.get_preferred_qonnx_opset()
):
def get_customop_wrapper(self, node, fallback_customop_version=util.get_preferred_qonnx_opset()):
"""Return CustomOp instance for given node, respecting the
imported opset version in the model protobuf. If the node's domain
is not found in the model's opset imports, fallback_customop_version
Expand Down
45 changes: 45 additions & 0 deletions tests/core/test_subgraph_traversal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,3 +281,48 @@ def test_analysis_fxn_without_apply_to_subgraphs_fails():
# Check that an analysis function fails when apply_to_subgraphs is False
model = make_subgraph_model(("top", [("sub1", []), ("sub2", [])]))
model.analysis(dummy_analysis_fxn, apply_to_subgraphs=True)


class SelfHandlingTransform(Transformation):
"""A transformation that manages its own subgraph traversal (e.g. because it
needs hierarchical context the generic descent cannot provide). It therefore
sets handles_subgraphs_internally = True and must not be driven via the
generic apply_to_subgraphs=True path."""

handles_subgraphs_internally = True

def __init__(self):
self.visited = list()

def apply(self, model_wrapper):
self.visited.append(model_wrapper.model.graph.name)
return model_wrapper, False


@pytest.mark.parametrize("cleanup", [False, True])
@pytest.mark.parametrize("make_deepcopy", [False, True])
@pytest.mark.parametrize(
"tree",
[
("top", []),
("top", [("sub1", []), ("sub2", [])]),
],
)
def test_self_handling_transform_rejects_apply_to_subgraphs(tree, cleanup, make_deepcopy):
# A transform that sets handles_subgraphs_internally = True must raise when
# driven via the generic apply_to_subgraphs=True path, regardless of the
# cleanup / make_deepcopy flags. The guard fires before any apply() runs.
model = make_subgraph_model(tree)
transform = SelfHandlingTransform()
with pytest.raises(ValueError, match="manages subgraph traversal itself"):
model.transform(transform, cleanup, make_deepcopy, apply_to_subgraphs=True)
assert transform.visited == []


def test_self_handling_transform_allows_top_level():
# Without apply_to_subgraphs the same transform runs normally (top-level only)
# and does not descend into subgraphs.
model = make_subgraph_model(("top", [("sub1", []), ("sub2", [])]))
transform = SelfHandlingTransform()
model.transform(transform, apply_to_subgraphs=False)
assert transform.visited == ["top"]
Loading