diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index a3d6075..6ae73e6 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -25,14 +25,14 @@ jobs: cd doc && python bug_summary.py - name: Test core run: | - pytest tests/core + pytest -x tests/core - name: Test PyTorch run: | pip install -r requirements/sys/torch.txt --pre --upgrade pip install -r requirements/sys/onnx.txt --pre --upgrade pip install -r requirements/sys/tvm.txt --pre --upgrade pip install -r requirements/sys/onnxruntime.txt --pre --upgrade - pytest tests/torch + pytest -x tests/torch yes | python nnsmith/cli/model_gen.py debug.viz=true model.type=torch mgen.method=symbolic yes | python nnsmith/cli/model_gen.py debug.viz=true model.type=torch mgen.method=symbolic-cinit yes | python nnsmith/cli/model_gen.py debug.viz=true model.type=torch backend.type="pt2 backend@inductor" mgen.method=concolic @@ -42,20 +42,20 @@ jobs: yes | python nnsmith/cli/model_gen.py model.type=torch mgen.method=symbolic-cinit mgen.rank_choices="[4]" mgen.dtype_choices="[f32]" mgen.include="[core.NCHWConv2d, core.ReLU]" mgen.patch_requires=./tests/mock/requires_patch.py backend.type=pt2 mgen.grad_check=true - name: Test ONNX + ONNXRuntime run: | - pytest tests/onnxruntime + pytest -x tests/onnxruntime yes | python nnsmith/cli/model_gen.py model.type=onnx mgen.method=symbolic yes | python nnsmith/cli/model_gen.py model.type=onnx backend.type=onnxruntime mgen.method=concolic python nnsmith/cli/model_exec.py model.type=onnx backend.type=onnxruntime model.path=nnsmith_output/model.onnx - name: Test ONNX + TVM run: | - pytest tests/tvm + pytest -x tests/tvm - name: Test ONNX + TRT run: | - pytest tests/tensorrt + pytest -x tests/tensorrt - name: Test TensorFlow run: | pip install -r requirements/sys/tensorflow.txt --pre --upgrade - pytest tests/tensorflow --log-cli-level=DEBUG + pytest -x tests/tensorflow --log-cli-level=DEBUG yes | python nnsmith/cli/model_gen.py model.type=tensorflow mgen.method=symbolic python nnsmith/cli/model_exec.py model.type=tensorflow backend.type=xla model.path=nnsmith_output/model/ yes | python nnsmith/cli/model_gen.py model.type=tensorflow mgen.method=concolic diff --git a/experiments/README.md b/experiments/README.md index f77e174..a8d2448 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -107,7 +107,7 @@ Next, we will run NNSmith to dump a bunch of random test-cases for an SUT (say P > > ```shell > # PyTorch -> pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch +> pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch > # TensorFlow > pip install tf-nightly > ``` diff --git a/nnsmith/materialize/onnx/__init__.py b/nnsmith/materialize/onnx/__init__.py index 5faf9bd..8a6b381 100644 --- a/nnsmith/materialize/onnx/__init__.py +++ b/nnsmith/materialize/onnx/__init__.py @@ -260,7 +260,7 @@ def load(cls, path: PathLike) -> "ONNXModel": # FIXME: missing key(s) in state_dict: "mlist.0.data", "mlist.1.data". if os.path.exists(torch_path): ret.with_torch = True - ret.torch_model = cls.PTType.load(torch_path) + ret.torch_model = cls.PTType.load(torch_path).torch_model ret.full_input_like = ret.torch_model.input_like ret.full_output_like = ret.torch_model.output_like diff --git a/nnsmith/materialize/torch/parse.py b/nnsmith/materialize/torch/parse.py index 3a074a9..dbc87ff 100644 --- a/nnsmith/materialize/torch/parse.py +++ b/nnsmith/materialize/torch/parse.py @@ -1,8 +1,7 @@ import operator -from typing import Any, Dict, List, Union, cast +from typing import Any, Dict, List, cast import torch -import torch._dynamo as dynamo import torch.fx as fx import torch.nn as nn import torch.utils._pytree as pytree @@ -22,27 +21,26 @@ def run_node(self, n: fx.node.Node) -> Any: def parse(model: nn.Module, *example_args: List[torch.Tensor]) -> GraphIR: - gm: fx.GraphModule = dynamo.export(model, *example_args)[0] + gm: fx.GraphModule = fx.symbolic_trace(model) # store shape info on nodes sp = PropInterpreter(gm) sp.run(*example_args) - def load_args(args: Union[List, Dict[str, Any]]) -> Union[List, Dict[str, Any]]: - """ - Map nodes to their outputs while keeping structures and other values the same. - """ - return torch.fx.graph.map_arg(args, lambda n: n.meta["res"]) - named_modules = dict(gm.named_modules()) ir = GraphIR() name2retvals: Dict[str, List[str]] = {} - for i_node, node in enumerate(gm.graph.nodes): + for node in gm.graph.nodes: node = cast(fx.node.Node, node) if node.op == "placeholder": - iexpr = InstExpr(Input(dim=len(node.meta["res"].shape)), []) + input_node = Input(dim=len(node.meta["res"].shape)) + input_node.abs_tensor = AbsTensor( + shape=list(node.meta["res"].shape), + dtype=DType.from_torch(node.meta["res"].dtype), + ) + iexpr = InstExpr(input_node, []) else: - args_flatten, args_treespec = pytree.tree_flatten(node.args) - kwargs_flatten, kwargs_treespec = pytree.tree_flatten(node.kwargs) + args_flatten, _ = pytree.tree_flatten(node.args) + kwargs_flatten, _ = pytree.tree_flatten(node.kwargs) input_nodes = [ a for a in (args_flatten + kwargs_flatten) @@ -67,8 +65,8 @@ def load_args(args: Union[List, Dict[str, Any]]) -> Union[List, Dict[str, Any]]: pytree.tree_flatten(node.meta["res"])[0], ) ) - nodes2empty = ( - lambda n: ConcreteOp.empty if isinstance(n, fx.node.Node) else n + nodes2empty = lambda n: ( + ConcreteOp.empty if isinstance(n, fx.node.Node) else n ) args_wo_nodes = pytree.tree_map(nodes2empty, node.args) kwargs_wo_nodes = pytree.tree_map(nodes2empty, node.kwargs) diff --git a/requirements/sys/torch.txt b/requirements/sys/torch.txt index cf4fc37..1dfee20 100644 --- a/requirements/sys/torch.txt +++ b/requirements/sys/torch.txt @@ -1,4 +1,4 @@ # TODO(@ganler): make other platform/device distribution also work. ---extra-index-url https://download.pytorch.org/whl/nightly/cpu +--index-url https://download.pytorch.org/whl/nightly/cpu --pre torch diff --git a/tests/torch/test_dump_load.py b/tests/torch/test_dump_load.py index 205637b..28b4773 100644 --- a/tests/torch/test_dump_load.py +++ b/tests/torch/test_dump_load.py @@ -45,7 +45,7 @@ def test_onnx_load_dump(tmp_path): # check oracle compare_two_oracle(oracle, loaded_testcase.oracle) - loaded_model = loaded_testcase.model.torch_model + loaded_model = loaded_testcase.model loaded_model.sat_inputs = {k: torch.from_numpy(v) for k, v in oracle.input.items()} rerun_oracle = loaded_model.make_oracle() compare_two_oracle(oracle, rerun_oracle) @@ -77,7 +77,7 @@ def test_bug_report_load_dump(tmp_path): # check oracle compare_two_oracle(oracle, loaded_testcase.oracle) - loaded_model = loaded_testcase.model.torch_model + loaded_model = loaded_testcase.model loaded_model.sat_inputs = {k: torch.from_numpy(v) for k, v in oracle.input.items()} rerun_oracle = loaded_model.make_oracle() compare_two_oracle(oracle, rerun_oracle)