├── nnsmith ├── cli │ ├── __init__.py │ ├── dtype_test.py │ ├── run_a_testcase.py │ ├── model_gen.py │ └── test_eq_helper.py ├── config │ ├── __init__.py │ ├── hydra │ │ └── job_logging │ │ │ ├── console.yaml │ │ │ └── file.yaml │ └── main.yaml ├── backends │ ├── __init__.py │ ├── hidet.py │ ├── xla.py │ ├── infinitensor.py │ ├── tensorrt.py │ ├── onnxruntime.py │ ├── pt2.py │ ├── torchjit.py │ ├── tvm.py │ └── tflite.py ├── abstract │ ├── __init__.py │ ├── extension.py │ └── tensor.py ├── __init__.py ├── macro.py ├── logging.py ├── difftest.py ├── materialize │ ├── torch │ │ ├── numeric.py │ │ ├── dialect.py │ │ ├── proxy_grad.py │ │ ├── input_gen.py │ │ └── parse.py │ └── tensorflow │ │ └── tfnet.py ├── error.py ├── filter.py └── util.py ├── requirements ├── sys │ ├── onnx.txt │ ├── tensorrt.txt │ ├── tensorflow.txt │ ├── onnxruntime.txt │ ├── tvm.txt │ └── torch.txt ├── exp.txt ├── dev.txt ├── core.txt └── install_all.sh ├── MANIFEST.in ├── env.sh ├── libdl_compiler_fuzzer_helper.so ├── exp_helper ├── 0_download_raw_data.sh ├── 3_eq_effectiveness_raw.py ├── tvm_log_process.py ├── evaluate_tvm_log.py ├── test_env.py ├── 1_sumarize_buglist_raw.py ├── process_diff_effect.py ├── 3_run_eq_effectiveness.py ├── 4_throughput_raw.py └── 4_run_throughput.py ├── pyproject.toml ├── tests ├── mock │ ├── requires_patch.py │ └── filter_patch.py ├── core │ ├── test_native_pickable.py │ ├── test_extra_constraint.py │ └── test_parse_name_kwargs.py ├── tensorflow │ ├── test_dump_load.py │ ├── test_tflite_backend.py │ └── test_xla_backend.py ├── onnxruntime │ └── test_ort_backend.py ├── tvm │ └── test_tvm_backend.py ├── torch │ ├── test_pt2_backend.py │ ├── test_torchjit_backend.py │ ├── test_biconvert.py │ └── test_dump_load.py └── tensorrt │ └── test_trt_backend.py ├── experiments ├── legacy │ ├── get_gentime_csv.py │ ├── README.md │ ├── input_search_exp.sh │ ├── cov_exp.sh │ ├── invalid_rate_torch_init.py │ ├── cnt_uniq_crash.py │ ├── plot_cov.py │ ├── compare_tzer.py │ ├── plot_inp_search.py │ ├── nnsmith_gen_onnx.py │ ├── lemon_tf2onnx.py │ ├── batch_eval.py │ └── plot_inp_search_merge.py └── evaluate_models.py ├── .pre-commit-config.yaml ├── equality_saturation_helper ├── rust_helper │ ├── src │ │ └── config.rs │ └── Cargo.toml ├── data_format.py └── fuzz_loop.py ├── doc ├── known-issues.md ├── log-and-err.md ├── concept.md ├── cli.md └── CONTRIBUTING.md ├── setup.cfg ├── .gitignore ├── readme.md └── CODE_OF_CONDUCT.md /nnsmith/cli/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /nnsmith/config/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /requirements/sys/onnx.txt: -------------------------------------------------------------------------------- 1 | onnx 2 | -------------------------------------------------------------------------------- /requirements/exp.txt: -------------------------------------------------------------------------------- 1 | pandas==1.3.5 2 | tqdm 3 | -------------------------------------------------------------------------------- /nnsmith/backends/__init__.py: -------------------------------------------------------------------------------- 1 | from nnsmith.backends.factory import BackendFactory 2 | -------------------------------------------------------------------------------- /requirements/dev.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | pre-commit 3 | black 4 | wget 5 | gputil 6 | pygraphviz 7 | -------------------------------------------------------------------------------- /nnsmith/abstract/__init__.py: -------------------------------------------------------------------------------- 1 | from .dtype import * 2 | from .op import * 3 | from .tensor import * 4 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | # include yaml files under nnsmith/config/**/*.yaml 2 | include nnsmith/config/**/*.yaml 3 | -------------------------------------------------------------------------------- /requirements/sys/tensorrt.txt: -------------------------------------------------------------------------------- 1 | --extra-index-url https://pypi.ngc.nvidia.com 2 | nvidia-pyindex 3 | nvidia-tensorrt 4 | -------------------------------------------------------------------------------- /env.sh: -------------------------------------------------------------------------------- 1 | export PYTHONPATH="/root/polyjuice" 2 | export EQ_HELPER_PATH="/root/polyjuice/libdl_compiler_fuzzer_helper.so" 3 | -------------------------------------------------------------------------------- /libdl_compiler_fuzzer_helper.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ChijinZ/PolyJuice-Fuzzer/HEAD/libdl_compiler_fuzzer_helper.so -------------------------------------------------------------------------------- /requirements/core.txt: -------------------------------------------------------------------------------- 1 | z3-solver>=4.11.0 2 | hydra-core>=1.2.0 3 | hydra-colorlog>=1.2.0 4 | multipledispatch 5 | appdirs 6 | numpy 7 | -------------------------------------------------------------------------------- /requirements/sys/tensorflow.txt: -------------------------------------------------------------------------------- 1 | tf-nightly 2 | # NOTE: https://www.tensorflow.org/install 3 | # https://www.tensorflow.org/install/pip 4 | -------------------------------------------------------------------------------- /nnsmith/__init__.py: -------------------------------------------------------------------------------- 1 | try: 2 | from nnsmith._version import __version__, __version_tuple__ 3 | except ImportError: 4 | __version__ = "local-dev" 5 | -------------------------------------------------------------------------------- /requirements/sys/onnxruntime.txt: -------------------------------------------------------------------------------- 1 | --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/ 2 | ort-nightly-gpu 3 | -------------------------------------------------------------------------------- /exp_helper/0_download_raw_data.sh: -------------------------------------------------------------------------------- 1 | cd /root/raw_data/ 2 | wget https://cloud.tsinghua.edu.cn/f/f1e2f3a09f6e4b88a59d/?dl=1 -O ./raw_data.tar.gz 3 | tar -xvf raw_data.tar.gz -------------------------------------------------------------------------------- /requirements/sys/tvm.txt: -------------------------------------------------------------------------------- 1 | # TODO(@ganler): switch back to pre-release after the merge 2 | # and release of https://github.com/apache/tvm/pull/13448 3 | apache-tvm==0.9.0 4 | -------------------------------------------------------------------------------- /requirements/sys/torch.txt: -------------------------------------------------------------------------------- 1 | # TODO(@ganler): make other platform/device distribution also work. 2 | --extra-index-url https://download.pytorch.org/whl/nightly/cpu 3 | --pre 4 | torch 5 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=45", "setuptools_scm[toml]>=6.2"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [tool.setuptools_scm] 6 | write_to = "nnsmith/_version.py" 7 | version_scheme = "release-branch-semver" 8 | local_scheme = "no-local-version" 9 | -------------------------------------------------------------------------------- /nnsmith/macro.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | ONNX_EXTERNAL_DATA_DIR_SUFFIX = "-mlist" 4 | NNSMITH_ORT_INTRA_OP_THREAD = int(os.getenv("NNSMITH_ORT_INTRA_OP_THREAD", 1)) 5 | NNSMITH_BUG_PATTERN_TOKEN = "${PATTERN}" 6 | 7 | 8 | def onnx2external_data_dir(onnx_file): 9 | return onnx_file + ONNX_EXTERNAL_DATA_DIR_SUFFIX 10 | -------------------------------------------------------------------------------- /tests/mock/requires_patch.py: -------------------------------------------------------------------------------- 1 | from nnsmith.abstract.arith import nnsmith_lt 2 | from nnsmith.abstract.extension import patch_requires 3 | 4 | 5 | @patch_requires("global", "core.NCHWConv2d") 6 | def limit_conv2d(self, _): 7 | # let the kernels to be > 3 8 | return [nnsmith_lt(3, self.kernel_h_size), nnsmith_lt(3, self.kernel_w_size)] 9 | -------------------------------------------------------------------------------- /tests/mock/filter_patch.py: -------------------------------------------------------------------------------- 1 | from nnsmith.filter import filter 2 | from nnsmith.materialize import BugReport 3 | 4 | 5 | @filter("test_fn") 6 | def filter_fn(report: BugReport): 7 | return False # Won't filter anything. 8 | 9 | 10 | @filter("test_cls") 11 | class FilterCls: 12 | def __call__(self, report: BugReport) -> bool: 13 | return False # Won't filter anything. 14 | -------------------------------------------------------------------------------- /experiments/legacy/get_gentime_csv.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import pandas as pd 3 | import os 4 | 5 | fuzz_report_rt = sys.argv[1] 6 | df = pd.read_pickle(os.path.join(fuzz_report_rt, "profile.pkl")) 7 | mstr = [f"{i}.onnx" for i in range(len(df))] 8 | df["model_name"] = mstr 9 | df["gen_t"] = df["gen_t"] 10 | df1 = df[["gen_t", "model_name"]] 11 | df1.to_csv(os.path.join(fuzz_report_rt, "gentime.csv"), index=False, header=False) 12 | -------------------------------------------------------------------------------- /nnsmith/backends/hidet.py: -------------------------------------------------------------------------------- 1 | from nnsmith.backends.pt2 import PT2 2 | 3 | 4 | class Hidet(PT2): 5 | def __init__(self, target: str = "cpu", optmax: bool = True, **kwargs): 6 | if target != "cuda": 7 | raise ValueError("Hidet backend only supports GPU!") 8 | super().__init__(target, optmax, backend="hidet", **kwargs) 9 | 10 | @property 11 | def system_name(self) -> str: 12 | return "hidet" 13 | -------------------------------------------------------------------------------- /exp_helper/3_eq_effectiveness_raw.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | from exp_helper.process_diff_effect import process_diff_effect 5 | 6 | RAW_DATA_PATH = "/root/raw_data/raw_data/tab4_raw_data/ablation_exp_5000.csv" 7 | OUTPUT_DIR_PATH = "/root/raw_data/raw_data_results" 8 | 9 | 10 | def main(): 11 | os.environ["OUTPUT_DIR"] = OUTPUT_DIR_PATH 12 | process_diff_effect(RAW_DATA_PATH) 13 | 14 | 15 | if __name__ == '__main__': 16 | main() 17 | -------------------------------------------------------------------------------- /nnsmith/config/hydra/job_logging/console.yaml: -------------------------------------------------------------------------------- 1 | version: 1 2 | formatters: 3 | colorlog: 4 | '()': 'colorlog.ColoredFormatter' 5 | format: "%(log_color)s%(levelname)-7s%(reset)s %(purple)s%(name)-6s%(reset)s - %(message)s" 6 | handlers: 7 | console: 8 | class: logging.StreamHandler 9 | formatter: colorlog 10 | stream: ext://sys.stdout 11 | root: 12 | level: INFO 13 | handlers: [console] 14 | 15 | disable_existing_loggers: false 16 | -------------------------------------------------------------------------------- /nnsmith/logging.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | VIZ_LOG = logging.getLogger("viz") 4 | FUZZ_LOG = logging.getLogger("fuzz") 5 | MGEN_LOG = logging.getLogger("mgen") 6 | SMT_LOG = logging.getLogger("smt") 7 | EXEC_LOG = logging.getLogger("exec") 8 | DTEST_LOG = logging.getLogger("dtest") 9 | CORE_LOG = logging.getLogger("core") 10 | HELPER_LOG = logging.getLogger("eq_helper") 11 | 12 | TF_LOG = logging.getLogger("gen|tf") 13 | TORCH_LOG = logging.getLogger("gen|torch") 14 | -------------------------------------------------------------------------------- /experiments/legacy/README.md: -------------------------------------------------------------------------------- 1 | # NOTES 2 | 3 | - Sripts under `experiments/legacy` are out of maintaince that it cannot work with the latest implementation. 4 | - For the current codebase, please refer to `experiments/README.md`. 5 | 6 | However, the legacy code might work with `620645967a14d6a7b077cedd9c2c03ed74af50d9` which is used in the ASPLOS'23 [artifact](http://nnsmith-asplos.rtfd.io/). 7 | For more information of using legacy code to render outputs in our artifact, please refer to bash scripts in the artifact [repository](https://github.com/ganler/nnsmith-asplos-artifact). 8 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pycqa/isort 3 | rev: 5.12.0 4 | hooks: 5 | - id: isort 6 | name: isort (python) 7 | args: ["--profile", "black"] 8 | exclude: | 9 | (?x)^( 10 | experiments/legacy/.*.py 11 | )$ 12 | - repo: https://github.com/psf/black 13 | rev: 22.6.0 14 | hooks: 15 | - id: black 16 | - repo: https://github.com/pre-commit/pre-commit-hooks 17 | rev: v4.3.0 18 | hooks: 19 | - id: check-yaml 20 | - id: end-of-file-fixer 21 | - id: trailing-whitespace 22 | -------------------------------------------------------------------------------- /experiments/legacy/input_search_exp.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export NMODELS=512 4 | export NNODES=${1:-20} # use `20` if the first argument is not given. 5 | export ROOT=$NMODELS-model-$NNODES-node-exp 6 | 7 | 8 | # Initial baseline. 9 | python experiments/input_search.py --max_nodes $NNODES --n_model $NMODELS --max_sample 1 --max_time_ms 8 --root $ROOT --result result-0.csv 10 | 11 | # 7 data points 12 | for i in {1..7} 13 | do 14 | echo "Running experiment $i" 15 | python experiments/input_search.py --max_nodes $NNODES --n_model $NMODELS --max_sample 1024 --max_time_ms $(($i * 8)) --root $ROOT --result result-$i.csv 16 | done 17 | -------------------------------------------------------------------------------- /equality_saturation_helper/rust_helper/src/config.rs: -------------------------------------------------------------------------------- 1 | use clap::Parser; 2 | 3 | /// Transfer an onnx model file to multiple semantically-equal onnx model files 4 | #[derive(Parser, Debug)] 5 | #[clap(author, version, about, long_about = None)] 6 | pub struct Args { 7 | #[clap(short, long, default_value = "100")] 8 | /// node limit of equality saturation 9 | pub node_limit: usize, 10 | #[clap(short, long, default_value = "5")] 11 | /// time limit (in second) of equality saturation 12 | pub time_limit_sec: u64, 13 | #[clap(short, long, default_value = "200")] 14 | /// iteration limit of equality saturation 15 | pub iter_limit: usize, 16 | } 17 | -------------------------------------------------------------------------------- /nnsmith/config/hydra/job_logging/file.yaml: -------------------------------------------------------------------------------- 1 | version: 1 2 | formatters: 3 | colorlog: 4 | '()': 'colorlog.ColoredFormatter' 5 | format: "%(log_color)s%(levelname)-7s%(reset)s %(purple)s%(name)-6s%(reset)s - %(message)s" 6 | simple: 7 | format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s' 8 | handlers: 9 | console: 10 | class: logging.StreamHandler 11 | formatter: colorlog 12 | stream: ext://sys.stdout 13 | file: 14 | class: logging.FileHandler 15 | formatter: simple 16 | filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log 17 | root: 18 | level: INFO 19 | handlers: [console, file] 20 | 21 | disable_existing_loggers: false 22 | -------------------------------------------------------------------------------- /exp_helper/tvm_log_process.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import re 3 | 4 | 5 | def main(): 6 | log_path = sys.argv[1] 7 | with open(log_path, 'r') as f: 8 | lines = f.readlines() 9 | 10 | for line in lines: 11 | if line.startswith("["): 12 | pattern = r'\[\d{2}:\d{2}:\d{2}\]\s' 13 | log_without_timestamp = re.sub(pattern, '', line) 14 | if "Cannot emit debug location for undefined span" not in log_without_timestamp \ 15 | and "Warning" not in log_without_timestamp \ 16 | and "naive_allocator" not in log_without_timestamp: 17 | print(log_without_timestamp, end='') 18 | 19 | 20 | if __name__ == '__main__': 21 | main() 22 | -------------------------------------------------------------------------------- /nnsmith/cli/dtype_test.py: -------------------------------------------------------------------------------- 1 | import hydra 2 | from omegaconf import DictConfig 3 | 4 | from nnsmith.backends import BackendFactory 5 | from nnsmith.materialize import Model 6 | from nnsmith.narrow_spec import auto_opconfig 7 | 8 | 9 | @hydra.main(version_base=None, config_path="../config", config_name="main") 10 | def main(cfg: DictConfig): 11 | backend_cfg = cfg["backend"] 12 | if backend_cfg["type"] is not None: 13 | factory = BackendFactory.init( 14 | name=backend_cfg["type"], 15 | target=backend_cfg["target"], 16 | optmax=backend_cfg["optmax"], 17 | parse_name=True, 18 | ) 19 | else: 20 | factory = None 21 | model_type = Model.init(cfg["model"]["type"], backend_target=backend_cfg["target"]) 22 | auto_opconfig(model_type, factory) 23 | 24 | 25 | if __name__ == "__main__": 26 | main() 27 | -------------------------------------------------------------------------------- /equality_saturation_helper/rust_helper/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "dl-compiler-fuzzer" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [lib] 9 | name = "dl_compiler_fuzzer_helper" 10 | #crate_type = ["staticlib", "rlib", "cdylib"] 11 | crate_type = ["cdylib"] 12 | path = "src/lib.rs" 13 | 14 | #[[bin]] 15 | #name = "dl_compiler_fuzzer" 16 | #path = "src/_main" 17 | 18 | [dependencies] 19 | egg = "0.9.4" 20 | itertools = "0.10.5" 21 | env_logger = "0.9" 22 | log = { version = "0.4", features = ["max_level_trace", "release_max_level_trace"] } 23 | serde_json = "1.0" 24 | bincode = "1.3" 25 | serde = { version = "1.0", features = ["derive"] } 26 | clap = { version = "3.1.18", features = ["derive"] } 27 | rand = "0.8" 28 | indexmap = "2.2.6" 29 | fxhash = "0.2.1" 30 | 31 | [profile.release] 32 | debug = true -------------------------------------------------------------------------------- /requirements/install_all.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -ex 4 | 5 | cd ./requirements 6 | 7 | # detect cuda 8 | # check if nvcc can work and return true or false 9 | check_nvcc() { 10 | if command -v nvcc >/dev/null; then 11 | return 0 12 | else 13 | return 1 14 | fi 15 | } 16 | 17 | pip install pip --upgrade # upgrade pip 18 | 19 | find . -name '*.txt' | while read -r dep_file; do 20 | if [ -f "$dep_file" ]; then 21 | # skip tensorrt if cuda is not available 22 | if [[ "$dep_file" == *"tensorrt"* ]] && ! check_nvcc; then 23 | echo "Skipping tensorrt" 24 | continue 25 | fi 26 | if [[ "$dep_file" == *"sys"* ]]; then 27 | # Files under sys should be nightly releases 28 | pip install -r "$dep_file" --upgrade --pre 29 | else 30 | pip install -r "$dep_file" 31 | fi 32 | fi 33 | done 34 | -------------------------------------------------------------------------------- /tests/core/test_native_pickable.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | 3 | import pytest 4 | 5 | from nnsmith.abstract.dtype import DType 6 | from nnsmith.abstract.op import MaxPool2d 7 | from nnsmith.abstract.tensor import AbsTensor 8 | 9 | 10 | def test_dtype_picklable(): 11 | dtype = DType.from_str("float32") 12 | dtype2 = pickle.loads(pickle.dumps(dtype)) 13 | assert dtype == dtype2 14 | 15 | 16 | def test_abstensor_picklable(): 17 | shape = [1, 2, 3] 18 | dtype = DType.from_str("float32") 19 | tensor = AbsTensor(shape, dtype) 20 | tensor2 = pickle.loads(pickle.dumps(tensor)) 21 | assert tensor == tensor2 22 | assert tensor.shape == tensor2.shape 23 | assert tensor.dtype == tensor2.dtype 24 | 25 | 26 | def test_absop_picklable(): 27 | maxpool = MaxPool2d(kh=2, kw=2, stride=1, padding=0) 28 | maxpool2 = pickle.loads(pickle.dumps(maxpool)) 29 | # TODO(@ganler): make AbsOp comparable. 30 | # assert maxpool == maxpool2 31 | -------------------------------------------------------------------------------- /doc/known-issues.md: -------------------------------------------------------------------------------- 1 | ## Incompatibility of TensorFlow-GPU over fork-based crash safty 2 | 3 | `fuzz.crash_safe=true` allows running compilation & execution in a forked process as a sandbox to catch crash and timeout. However, CUDA runtime is not compatible with fork. In tensorflow, the symptom is crash in forked subprocess: 4 | 5 | ```txt 6 | F tensorflow/stream_executor/cuda/cuda_driver.cc:219] Failed setting context: CUDA_ERROR_NOT_INITIALIZED: initialization error 7 | ``` 8 | 9 | - For `tflite` it's okay as it does not require GPU and `nnsmith.fuzz` will directly set `CUDA_VISIBLE_DEVICES=-1` in the beginning; 10 | - For `xla` it's a bit headache, currently we need to manually specify `fuzz.crash_safe=false` for fuzzing and allow it to crash; 11 | - We are tracking this [issue](https://github.com/tensorflow/tensorflow/issues/57877) in TensorFlow. We are likely to fix this by executing a TensorFlow model in a seperated process if it cannot be resolved in the near future. 12 | -------------------------------------------------------------------------------- /nnsmith/abstract/extension.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Type 2 | 3 | REQUIRES_PATCH = {} 4 | ACTIVATED_PATCH = {} 5 | 6 | 7 | class patch_requires: 8 | def __init__(self, tag: str, opname: str): 9 | self.tag = tag 10 | self.opname = opname 11 | 12 | def __call__(self, f): 13 | REQUIRES_PATCH.setdefault(self.tag, {}).setdefault(self.opname, []).append(f) 14 | return f 15 | 16 | 17 | def activate_ext( 18 | opset: List[Type["AbsOpBase"]], factory: Optional["BackendFactory"] = None 19 | ): 20 | for op in opset: 21 | if "global" in REQUIRES_PATCH: 22 | ACTIVATED_PATCH.setdefault(op.name(), []).extend( 23 | REQUIRES_PATCH["global"].get(op.name(), []) 24 | ) 25 | 26 | if factory is not None and factory.system_name in REQUIRES_PATCH: 27 | ACTIVATED_PATCH.setdefault(op.name(), []).extend( 28 | REQUIRES_PATCH[factory.system_name].get(op.name(), []) 29 | ) 30 | -------------------------------------------------------------------------------- /exp_helper/evaluate_tvm_log.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import difflib 3 | 4 | 5 | # return the difference rate of two log files 6 | def evaluate(log1_lines, log2_lines) -> float: 7 | # Determine which log is longer 8 | long_log, short_log = (log1_lines, log2_lines) if len(log1_lines) > len(log2_lines) else (log2_lines, log1_lines) 9 | 10 | diff = difflib.unified_diff(long_log, short_log) 11 | 12 | # diff_lines = len([l for l in diff if l.startswith('-')]) 13 | diff_lines = len(list(diff)) 14 | 15 | return diff_lines / len(long_log) 16 | 17 | 18 | def main(): 19 | assert len(sys.argv) == 3, "Usage: python evaluate_tvm_log.py log_path_1 log_path_2" 20 | 21 | log_path_1 = sys.argv[1] 22 | log_path_2 = sys.argv[2] 23 | 24 | with open(log_path_1, 'r') as f: 25 | lines_1 = f.readlines() 26 | 27 | with open(log_path_2, 'r') as f: 28 | lines_2 = f.readlines() 29 | 30 | print(evaluate(lines_1, lines_2)) 31 | 32 | 33 | if __name__ == '__main__': 34 | main() 35 | -------------------------------------------------------------------------------- /nnsmith/cli/run_a_testcase.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | from equivalent_fuzz import EquivalentFuzzingLoop, TestCase, BugReport 5 | import pickle 6 | 7 | 8 | def main(): 9 | fuzzer_object_path = sys.argv[1] 10 | test_case_path = sys.argv[2] 11 | 12 | tmp_file_path = None 13 | 14 | if len(sys.argv) == 4: 15 | tmp_file_path = sys.argv[3] 16 | 17 | assert os.path.exists(fuzzer_object_path), f"{fuzzer_object_path} does not exist" 18 | 19 | with open(fuzzer_object_path, "rb") as f: 20 | fuzzer_object: EquivalentFuzzingLoop = pickle.load(f) 21 | 22 | model_type = fuzzer_object.ModelType 23 | testcase = TestCase.load(model_type, test_case_path) 24 | 25 | if tmp_file_path is None: 26 | exit(101) 27 | 28 | res = fuzzer_object.execute_testcase(testcase, crash_safe=True, file_path_for_subprocess_log=tmp_file_path, 29 | timeout=20) 30 | if isinstance(res, BugReport): 31 | exit(100) 32 | 33 | 34 | if __name__ == '__main__': 35 | main() 36 | -------------------------------------------------------------------------------- /exp_helper/test_env.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | import subprocess 3 | import os 4 | import re 5 | 6 | 7 | def main(): 8 | print("start to test environment") 9 | flag = True 10 | cpu_count = multiprocessing.cpu_count() 11 | if cpu_count < 16: 12 | print( 13 | f"WARNING: your cpu count is {cpu_count}, this may not satisfy the requirement. " 14 | f"You can use raw-data for reproducing the results. From-scratch reproduction may encounter some problems") 15 | flag = False 16 | meminfo = dict((i.split()[0].rstrip(':'), int(i.split()[1])) for i in open('/proc/meminfo').readlines()) 17 | mem_Gib = meminfo['MemTotal'] / 1024 / 1024 18 | if mem_Gib < 32: 19 | print(f"WARNING: your total memory size is {mem_Gib} GB, this may not satisfy the requirement" 20 | f"You can use raw-data for reproducing the results. From-scratch reproduction may encounter some problems") 21 | flag = False 22 | 23 | if flag: 24 | print("Your environment is suitable for reproducing the results") 25 | 26 | 27 | if __name__ == "__main__": 28 | main() 29 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = nnsmith 3 | description = "Automatic DNN generation for fuzzing and more." 4 | long_description = file: README.md 5 | long_description_content_type = text/markdown 6 | url = https://github.com/ise-uiuc/nnsmith 7 | license = Apache-2.0 8 | license_file = LICENSE 9 | platform = any 10 | 11 | [options] 12 | packages = find: 13 | python_requires = >=3.8 14 | dependency_links = 15 | install_requires = 16 | z3-solver>=4.11.0 17 | hydra-core>=1.2.0 18 | hydra-colorlog>=1.2.0 19 | multipledispatch>=0.6.0 20 | appdirs>=1.4.4 21 | numpy 22 | 23 | # TODO: make it nightly. 24 | [options.extras_require] 25 | onnx = torch 26 | onnx 27 | onnxruntime = onnxruntime 28 | torch 29 | onnx 30 | tensorflow = tf-nightly 31 | torch = torch 32 | tvm = apache-tvm 33 | torch 34 | onnx 35 | 36 | [options.package_data] 37 | nnsmith = config/**/*.yaml 38 | 39 | [options.entry_points] 40 | console_scripts = 41 | nnsmith.model_gen = nnsmith.cli.model_gen:main 42 | nnsmith.model_exec = nnsmith.cli.model_exec:main 43 | nnsmith.dtype_test = nnsmith.cli.dtype_test:main 44 | nnsmith.fuzz = nnsmith.cli.fuzz:main 45 | -------------------------------------------------------------------------------- /equality_saturation_helper/data_format.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import os 3 | import torch 4 | 5 | root_path = "/workspace/modified_nnsmith/triage_bugs/infini_11_27/fuzz_report_cpu_11_27/cpu_transpose" 6 | torch_model = os.path.join(root_path, "t1", "model.pth") 7 | data_file = os.path.join(root_path, "0.pickle") 8 | 9 | idx = 0 10 | key_list = [] 11 | print(pickle.load(open(data_file, 'rb'))) 12 | for k, v in pickle.load(open(data_file, 'rb')).items(): 13 | print(f"data_{idx} = np.random.normal(5, 1, size={v.shape}).astype(np.{v.dtype})") 14 | key_list.append(k) 15 | idx += 1 16 | 17 | # dict_str = "input_data_0 = [" 18 | # for i in range(idx): 19 | # dict_str += f"data_{i}," 20 | # dict_str += "]" 21 | # print(dict_str) 22 | 23 | dict_str = "input_dict_0 = {" 24 | for i in range(idx): 25 | dict_str += f"'{key_list[i]}': data_{i}, " 26 | dict_str += "}" 27 | print(dict_str) 28 | 29 | 30 | # print("\n") 31 | # model_consts = torch.load(torch_model) 32 | # p_idx = 0 33 | # for k, v in model_consts.items(): 34 | # if sum(list(v.shape)) < 10: 35 | # print(v) 36 | # print(f"p{p_idx} = np.random.normal(0, 3, size={tuple(v.shape)}).astype(np.{v.dtype})") 37 | # print(f"p{p_idx} = torch.from_numpy(p{p_idx}).to(DEVICE)") -------------------------------------------------------------------------------- /experiments/legacy/cov_exp.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export NNSMITH_DCE=0.1 3 | # TVM. 4 | export LIB_PATH='../tvm/build/libtvm.so ../tvm/build/libtvm_runtime.so' 5 | start_time=`date +%s` 6 | python nnsmith/fuzz.py --time 14400 --max_nodes 10 --eval_freq 256 \ 7 | --mode random --backend tvm --root nnsmith-tvm-base 8 | exp0_t=$(expr `date +%s` - $start_time) 9 | 10 | start_time=`date +%s` 11 | python nnsmith/fuzz.py --time 14400 --max_nodes 10 --eval_freq 256 \ 12 | --mode guided --backend tvm --root nnsmith-tvm-guided 13 | exp1_t=$(expr `date +%s` - $start_time) 14 | 15 | # ONNXRuntime. 16 | export LIB_PATH='../onnxruntime/build/Linux/RelWithDebInfo/libonnxruntime_providers_shared.so ../onnxruntime/build/Linux/RelWithDebInfo/libonnxruntime.so' 17 | start_time=`date +%s` 18 | python nnsmith/fuzz.py --time 14400 --max_nodes 10 --eval_freq 256 \ 19 | --mode random --backend ort --root nnsmith-ort-base 20 | exp2_t=$(expr `date +%s` - $start_time) 21 | 22 | start_time=`date +%s` 23 | python nnsmith/fuzz.py --time 14400 --max_nodes 10 --eval_freq 256 \ 24 | --mode guided --backend ort --root nnsmith-ort-guided 25 | exp3_t=$(expr `date +%s` - $start_time) 26 | 27 | echo "Experiment time of last 4 runs: '$exp0_t','$exp1_t','$exp2_t','$exp3_t' seconds." 28 | -------------------------------------------------------------------------------- /experiments/legacy/invalid_rate_torch_init.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import random 3 | 4 | from tqdm import tqdm 5 | 6 | from nnsmith.graph_gen import random_model_gen, SymbolNet 7 | 8 | 9 | def mknet(args): 10 | model_seed = random.getrandbits(32) 11 | gen, solution = random_model_gen( 12 | mode=args.mode, seed=model_seed, max_nodes=args.max_nodes, init_fp=True 13 | ) 14 | net = SymbolNet(gen.abstract_graph, solution, alive_shapes=gen.alive_shapes) 15 | net.eval() 16 | return net, gen.num_op(), model_seed 17 | 18 | 19 | if __name__ == "__main__": 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("--max_nodes", type=int, default=20) 22 | parser.add_argument("--n_model", type=int, default=1000) 23 | parser.add_argument("--mode", type=str, default="random") 24 | args = parser.parse_args() 25 | 26 | n_invalid = 0 27 | with tqdm(range(args.n_model)) as pbar: 28 | for model_id in pbar: 29 | net, num_op, model_seed = mknet(args) 30 | net.check_intermediate_numeric = True 31 | _ = net(*net.get_random_inps(base="center", margin=1)) 32 | n_invalid += net.invalid_found_last 33 | cur_model_size = model_id + 1 34 | pbar.set_description( 35 | f"invalid:{n_invalid}/{cur_model_size} = {100 * n_invalid/cur_model_size:.2f}%" 36 | ) 37 | -------------------------------------------------------------------------------- /tests/core/test_extra_constraint.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from nnsmith.abstract.arith import nnsmith_lt 4 | from nnsmith.abstract.extension import activate_ext, patch_requires 5 | from nnsmith.abstract.op import Constant, Input, NCHWConv2d, ReLU 6 | from nnsmith.graph_gen import model_gen 7 | 8 | 9 | def test_only_conv_relu(): 10 | gen = model_gen( 11 | opset=[ReLU, NCHWConv2d], 12 | max_nodes=5, 13 | rank_choices=(4,), 14 | dtype_choices=("float32",), 15 | ) 16 | 17 | ir = gen.make_concrete() 18 | 19 | for inst in ir.insts: 20 | assert type(inst.iexpr.op) in [ReLU, Input, Constant, NCHWConv2d] 21 | 22 | 23 | def test_constrain_conv_ksize(): 24 | @patch_requires("global", "core.NCHWConv2d") 25 | def limit_conv2d(self, _): 26 | # let the kernels to be > 3 27 | return [nnsmith_lt(3, self.kernel_h_size), nnsmith_lt(3, self.kernel_w_size)] 28 | 29 | opset = [ReLU, NCHWConv2d] 30 | activate_ext(opset) 31 | gen = model_gen( 32 | opset=opset, 33 | max_nodes=5, 34 | rank_choices=(4,), 35 | dtype_choices=("float32",), 36 | ) 37 | 38 | ir = gen.make_concrete() 39 | for inst in ir.insts: 40 | assert type(inst.iexpr.op) in [ReLU, Input, Constant, NCHWConv2d] 41 | if isinstance(inst.iexpr.op, NCHWConv2d): 42 | assert inst.iexpr.op.kernel_h_size > 3 43 | assert inst.iexpr.op.kernel_w_size > 3 44 | -------------------------------------------------------------------------------- /doc/log-and-err.md: -------------------------------------------------------------------------------- 1 | ## Logging 2 | 3 | ### Modularization 4 | 5 | We support the following logging "keys": 6 | 7 | - `fuzz`: fuzzing loop; 8 | - `mgen`: model generation; 9 | - `smt`: constraints in smt solving; 10 | - `exec`: model execution; 11 | - `viz`: graphviz visualization; 12 | - `dtest`: dtype_test; 13 | - `core`: seed setting, etc; 14 | 15 | The show messages above "INFO" level (see [Python's logging module](https://docs.python.org/3/library/logging.html)). To show debug level message, add `hydra.verbose=[${keys}]` (also see [hydra.logging](https://hydra.cc/docs/1.2/tutorials/basic/running_your_app/logging/)). 16 | 17 | ```shell 18 | # Show debug information related to `fuzz`: 19 | ${NNSMITH_CMD} hydra.verbose=fuzz 20 | # Show debug info for `fuzz` and `exec`: 21 | ${NNSMITH_CMD} hydra.verbose="[fuzz,exec]" 22 | ``` 23 | 24 | #### Where the log is? 25 | 26 | By default, NNSmith logs things both in `console` and `file`. You can find the loggings in [`outputs/${DATE}/${APP}.log`](https://hydra.cc/docs/1.2/tutorials/basic/running_your_app/working_directory/) (current working directory). 27 | 28 | ## Errors 29 | 30 | See `nnsmith/error.py`: 31 | 32 | - `ConstraintError`: Unsatisfiable constraints which is a hint to re-try; 33 | - `InternalError`: NNSmith has some internal bugs that should be fixed. 34 | 35 | Takeways: 36 | 37 | - Catch `ConstraintError` as a hint to re-try graph generation (no satisfiable solution yet); 38 | - Never catch `InternalError` -- but let the maintainer know the issue and fix it. 39 | -------------------------------------------------------------------------------- /nnsmith/difftest.py: -------------------------------------------------------------------------------- 1 | from typing import Dict 2 | 3 | import numpy as np 4 | from numpy import testing 5 | 6 | 7 | def assert_allclose( 8 | actual: Dict[str, np.ndarray], 9 | desired: Dict[str, np.ndarray], 10 | actual_name: str, 11 | oracle_name: str, 12 | equal_nan=False, 13 | rtol=1e-2, 14 | atol=1e-3, 15 | ): 16 | akeys = set(actual.keys()) 17 | dkeys = set(desired.keys()) 18 | if akeys != dkeys: 19 | raise KeyError(f"{actual_name}: {akeys} != {oracle_name}: {dkeys}") 20 | 21 | for key in akeys: 22 | lhs = actual[key] 23 | rhs = desired[key] 24 | 25 | if lhs is not None and rhs is not None: 26 | # check if lhs is np.ndarray 27 | if lhs is not None and not isinstance(lhs, np.ndarray): 28 | raise TypeError( 29 | f"{actual_name}[{key}] is not np.ndarray but {type(lhs)}" 30 | ) 31 | 32 | # check if rhs is np.ndarray 33 | if rhs is not None and not isinstance(rhs, np.ndarray): 34 | raise TypeError( 35 | f"{oracle_name}[{key}] is not np.ndarray but {type(rhs)}" 36 | ) 37 | 38 | testing.assert_allclose( 39 | lhs, 40 | rhs, 41 | equal_nan=equal_nan, 42 | rtol=rtol, 43 | atol=atol, 44 | err_msg=f"{actual_name} != {oracle_name} at {key}", 45 | ) 46 | else: 47 | return lhs is None and rhs is None 48 | -------------------------------------------------------------------------------- /tests/tensorflow/test_dump_load.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | 4 | from nnsmith.graph_gen import model_gen 5 | from nnsmith.materialize import TestCase 6 | from nnsmith.materialize.tensorflow import TFModelCPU, tf_dict_from_np 7 | from nnsmith.narrow_spec import auto_opset 8 | 9 | TestCase.__test__ = False # supress PyTest warning 10 | 11 | 12 | def test_onnx_load_dump(tmp_path): 13 | d = tmp_path / "test_onnx_load_dump" 14 | d.mkdir() 15 | 16 | gen = model_gen( 17 | opset=auto_opset(TFModelCPU), 18 | seed=54341, 19 | max_nodes=5, 20 | ) 21 | 22 | ir = gen.make_concrete() 23 | 24 | model = TFModelCPU.from_gir(ir) 25 | 26 | model.refine_weights() # either random generated or gradient-based. 27 | oracle = model.make_oracle() 28 | 29 | testcase = TestCase(model, oracle) 30 | testcase.dump(root_folder=d) 31 | 32 | loaded_testcase = TestCase.load(model_type=type(model), root_folder=d) 33 | 34 | def compare_two_oracle(src, loaded): 35 | assert len(loaded.input) == len(src.input) 36 | assert len(loaded.output) == len(src.output) 37 | for k, v in loaded.input.items(): 38 | assert np.allclose(v, src.input[k], equal_nan=True) 39 | for k, v in loaded.output.items(): 40 | assert np.allclose(v, src.output[k], equal_nan=True) 41 | 42 | # check oracle 43 | compare_two_oracle(oracle, loaded_testcase.oracle) 44 | 45 | loaded_model = loaded_testcase.model 46 | rerun_oracle = loaded_model.make_oracle(tf_dict_from_np(oracle.input)) 47 | compare_two_oracle(oracle, rerun_oracle) 48 | -------------------------------------------------------------------------------- /nnsmith/backends/xla.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, List 2 | 3 | import tensorflow as tf # type: ignore 4 | from multipledispatch import dispatch 5 | 6 | from nnsmith.backends.factory import BackendCallable, BackendFactory 7 | from nnsmith.materialize.tensorflow import ( 8 | EagerModeCtx, 9 | TFModel, 10 | np_dict_from_tf, 11 | tf_dict_from_np, 12 | get_cuda_string 13 | ) 14 | 15 | 16 | class XLA(BackendFactory): 17 | def __init__(self, target="cpu", optmax: bool = True): 18 | super().__init__(target, optmax) 19 | 20 | if self.target == "cpu": 21 | self.device = tf.device(tf.config.list_logical_devices("CPU")[0].name) 22 | elif self.target == "cuda": 23 | self.device = tf.device(get_cuda_string()) 24 | else: 25 | raise ValueError( 26 | f"Unknown device: {self.target}. Only `cpu` and `cuda` are supported." 27 | ) 28 | 29 | @property 30 | def system_name(self) -> str: 31 | return "xla" 32 | 33 | @property 34 | def version(self) -> str: 35 | return tf.__version__ 36 | 37 | @property 38 | def import_libs(self) -> List[str]: 39 | return ["import tensorflow as tf"] 40 | 41 | @dispatch(TFModel) 42 | def make_backend(self, model: TFModel) -> BackendCallable: 43 | with self.device, EagerModeCtx(False): 44 | compiled = tf.function(jit_compile=True)(model.concrete_net()) 45 | 46 | def closure(inputs: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]: 47 | with self.device, EagerModeCtx(False): 48 | result = np_dict_from_tf(compiled(**tf_dict_from_np(inputs))) 49 | return result 50 | 51 | return closure 52 | -------------------------------------------------------------------------------- /experiments/legacy/cnt_uniq_crash.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from nnsmith.backends import mk_factory, BackendFactory 4 | 5 | if __name__ == "__main__": 6 | import argparse 7 | 8 | parser = argparse.ArgumentParser() 9 | parser.add_argument( 10 | "--backend", type=str, help="One of ort, trt, tvm, and xla", required=True 11 | ) 12 | parser.add_argument("--device", type=str, default="cpu") 13 | parser.add_argument("--root", type=str, help="Path to the bug folder") 14 | args = parser.parse_args() 15 | 16 | fac = mk_factory(args.backend, args.device, optmax=True) 17 | 18 | crash_msg = set() 19 | 20 | for dir in os.listdir(args.root): 21 | path = os.path.join(args.root, dir) 22 | if not os.path.isdir(path): 23 | continue 24 | if not dir.startswith("bug-"): 25 | continue 26 | if dir.startswith("bug-torch") or dir.startswith("bug-omin"): 27 | continue 28 | 29 | try: 30 | onnx_path = os.path.join(path, "model.onnx") 31 | onnx_model = BackendFactory.get_onnx_proto(onnx_path) 32 | try: 33 | fac.make_backend(onnx_model) 34 | except Exception as e: 35 | crash_msg.add(str(e)) 36 | except Exception as e: 37 | print(e) # filter model-too-large kind of things. 38 | continue 39 | 40 | print(f"{args.root} got {len(crash_msg)} different crash messages:") 41 | err_msg_path = os.path.join(args.root, "crash_msg.txt") 42 | print(f"Writing unique crash messages to {err_msg_path}") 43 | if crash_msg: 44 | with open(err_msg_path, "w") as f: 45 | for msg in crash_msg: 46 | print(msg, file=f) 47 | print("$\n", file=f) # splitter 48 | -------------------------------------------------------------------------------- /nnsmith/config/main.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | type: null 3 | path: "???" # can be multiple files tho. 4 | 5 | mgen: # model gen. 6 | max_nodes: 5 7 | timeout_ms: 10000 8 | vulops: False 9 | method: "symbolic-cinit" 10 | save: "nnsmith_output" 11 | seed: null 12 | max_elem_per_tensor: 65536 # 2^16 13 | rank_choices: null # 0 ~ __MAX_RANK__ 14 | dtype_choices: null # 0 ~ __MAX_DTYPE__ 15 | include: null # ops to include; example mgen.include="[core.NCHWConv2d, core.ReLU]" 16 | exclude: null # ops to exclude; 17 | patch_requires: [] # files that with @patch_requires 18 | grad_check: false # additionally check gradients 19 | 20 | # backend config 21 | backend: 22 | type: null 23 | optmax: true 24 | target: "cpu" 25 | 26 | ad: 27 | type: null 28 | 29 | cache: 30 | topset: true # Run dtype test with automatically maintained cache 31 | 32 | debug: 33 | viz: false 34 | viz_fmt: "png" # or "svg" for much smaller figure size and precision; 35 | 36 | fuzz: 37 | time: 14400 38 | root: "???" 39 | seed: null 40 | crash_safe: false 41 | test_timeout: null 42 | save_test: null 43 | random_num: 1 44 | 45 | filter: 46 | type: [] 47 | patch: [] 48 | 49 | cmp: 50 | equal_nan: true # skip regarding it as a bug if with fp exception values. 51 | 52 | raw_input: null # path to raw input data (Dict[str, np.ndarray]) 53 | 54 | oracle: "auto" 55 | # "auto": use `oracle.pkl` in local path; 56 | # PathLike: get the oracle from somewhere else; 57 | # null: fallback to random. 58 | 59 | with: 60 | type: null 61 | optmax: true 62 | target: "cpu" 63 | 64 | seed: null 65 | bug_presence: "report" # or "crash" 66 | save: null # path to save the bug report if `bug_presence` is "report" 67 | 68 | defaults: 69 | - override hydra/job_logging: file 70 | - override hydra/hydra_logging: colorlog 71 | -------------------------------------------------------------------------------- /tests/onnxruntime/test_ort_backend.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from nnsmith.abstract.dtype import DType 4 | from nnsmith.backends import BackendFactory 5 | from nnsmith.graph_gen import model_gen 6 | from nnsmith.materialize import Model, TestCase 7 | from nnsmith.narrow_spec import auto_opconfig, auto_opset 8 | 9 | TestCase.__test__ = False # supress PyTest warning 10 | 11 | 12 | def test_narrow_spec_cache_make_and_reload(): 13 | factory = BackendFactory.init("onnxruntime", target="cpu", optmax=True) 14 | ONNXModel = Model.init("onnx") 15 | opset_lhs = auto_opconfig(ONNXModel, factory) 16 | assert opset_lhs, "Should not be empty... Something must go wrong." 17 | opset_rhs = auto_opconfig(ONNXModel, factory) 18 | assert opset_lhs == opset_rhs 19 | 20 | # Assert types 21 | assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType) 22 | 23 | # Assert Dictionary Type Equality 24 | assert type(opset_lhs) == type(opset_rhs) 25 | assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"]) 26 | assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type( 27 | opset_rhs["core.ReLU"].in_dtypes[0][0] 28 | ) 29 | 30 | 31 | def test_synthesized_onnx_model(tmp_path): 32 | d = tmp_path / "test_ort_onnx" 33 | d.mkdir() 34 | 35 | ONNXModel = Model.init("onnx") 36 | 37 | factory = BackendFactory.init("onnxruntime", target="cpu", optmax=False) 38 | gen = model_gen( 39 | opset=auto_opset(ONNXModel, factory), 40 | seed=23132, 41 | max_nodes=2, 42 | ) 43 | 44 | model = ONNXModel.from_gir(gen.make_concrete()) 45 | 46 | assert model.with_torch 47 | 48 | model.refine_weights() # either random generated or gradient-based. 49 | oracle = model.make_oracle() 50 | 51 | testcase = TestCase(model, oracle) 52 | testcase.dump(root_folder=d) 53 | 54 | assert factory.verify_testcase(testcase) is None 55 | -------------------------------------------------------------------------------- /tests/tensorflow/test_tflite_backend.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from nnsmith.abstract.dtype import DType 4 | from nnsmith.backends import BackendFactory 5 | from nnsmith.graph_gen import model_gen 6 | from nnsmith.materialize import Model, TestCase 7 | from nnsmith.materialize.tensorflow import TFModelCPU 8 | from nnsmith.narrow_spec import auto_opconfig, auto_opset 9 | 10 | TestCase.__test__ = False # supress PyTest warning 11 | 12 | 13 | def test_narrow_spec_cache_make_and_reload(): 14 | factory = BackendFactory.init("tflite", target="cpu", optmax=True) 15 | ModelType = Model.init("tensorflow") 16 | opset_lhs = auto_opconfig(ModelType, factory) 17 | assert opset_lhs, "Should not be empty... Something must go wrong." 18 | opset_rhs = auto_opconfig(ModelType, factory) 19 | assert opset_lhs == opset_rhs 20 | 21 | # Assert types 22 | assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType) 23 | 24 | # Assert Dictionary Type Equality 25 | assert type(opset_lhs) == type(opset_rhs) 26 | assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"]) 27 | assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type( 28 | opset_rhs["core.ReLU"].in_dtypes[0][0] 29 | ) 30 | 31 | 32 | def test_synthesized_model(tmp_path): 33 | d = tmp_path / "test_tflite" 34 | d.mkdir() 35 | 36 | ModelType = Model.init("tensorflow") 37 | factory = BackendFactory.init("tflite", target="cpu", optmax=False) 38 | 39 | gen = model_gen( 40 | opset=auto_opset(TFModelCPU, factory), 41 | seed=23132, 42 | max_nodes=4, 43 | ) # One op should not be easily wrong... I guess. 44 | 45 | model = ModelType.from_gir(gen.make_concrete()) 46 | 47 | # model.refine_weights() # either random generated or gradient-based. 48 | oracle = model.make_oracle() 49 | 50 | testcase = TestCase(model, oracle) 51 | testcase.dump(root_folder=d) 52 | 53 | assert factory.verify_testcase(testcase) is None 54 | -------------------------------------------------------------------------------- /nnsmith/materialize/torch/numeric.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from multipledispatch import dispatch 3 | 4 | from nnsmith.abstract.op import * 5 | 6 | 7 | @torch.jit.ignore 8 | def numeric_valid(outputs) -> bool: 9 | with torch.no_grad(): 10 | return all([torch.isfinite(out).all() for out in outputs]) 11 | 12 | 13 | # generalized loss fn 14 | def smoothed_relu(x): 15 | if x.dtype == torch.float16: 16 | return torch.relu(x.float()).half() 17 | return torch.relu(x) 18 | 19 | 20 | def loss_ge_zero(x): 21 | return smoothed_relu(-x) 22 | 23 | 24 | def loss_le_zero(x): 25 | return smoothed_relu(x) 26 | 27 | 28 | def loss_lt_zero(x): 29 | return loss_le(x, -1e-10) 30 | 31 | 32 | def loss_gt_zero(x): 33 | return loss_ge(x, 1e-10) 34 | 35 | 36 | def loss_ge(x, y): 37 | return loss_ge_zero(x - y) 38 | 39 | 40 | def loss_le(x, y): 41 | return loss_le_zero(x - y) 42 | 43 | 44 | def loss_gt(x, y): 45 | return loss_gt_zero(x - y) 46 | 47 | 48 | def loss_lt(x, y): 49 | return loss_lt_zero(x - y) 50 | 51 | 52 | # loss_fn: backward 53 | 54 | 55 | @dispatch(Div) 56 | def loss_fn(op: Div): 57 | return lambda op, y: loss_gt_zero(torch.abs(y)) 58 | 59 | 60 | @dispatch(Pow) 61 | def loss_fn(op: Pow): 62 | def torch_loss(a, b): 63 | # a >= 0 && b*log(a) <= 20 64 | l0 = loss_gt_zero(a) 65 | if torch.any(l0 > 0): 66 | return ("l0", l0) 67 | l1 = loss_le( 68 | b * torch.log(torch.maximum(a, torch.tensor(1e-40, dtype=a.dtype))), 40 69 | ) 70 | return ("l1", l1) 71 | 72 | return torch_loss 73 | 74 | 75 | @dispatch(Acos) 76 | def loss_fn(op: Acos): 77 | return lambda x: loss_le(x.abs(), 1) 78 | 79 | 80 | @dispatch(Sqrt) 81 | def loss_fn(op: Sqrt): 82 | return lambda x: loss_ge(x, 0) 83 | 84 | 85 | @dispatch(Asin) 86 | def loss_fn(op: Asin): 87 | return lambda x: loss_le(x.abs(), 1) 88 | 89 | 90 | @dispatch(Log2) 91 | def loss_fn(op: Log2): 92 | return lambda x: loss_gt_zero(x) 93 | -------------------------------------------------------------------------------- /tests/tvm/test_tvm_backend.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import tvm 3 | 4 | from nnsmith.abstract.dtype import DType 5 | from nnsmith.backends import BackendFactory 6 | from nnsmith.graph_gen import model_gen 7 | from nnsmith.materialize import Model, TestCase 8 | from nnsmith.narrow_spec import auto_opconfig, auto_opset 9 | 10 | TestCase.__test__ = False # supress PyTest warning 11 | 12 | 13 | def test_narrow_spec_cache_make_and_reload(): 14 | factory = BackendFactory.init("tvm", target="cpu", optmax=True) 15 | ONNXModel = Model.init("onnx") 16 | opset_lhs = auto_opconfig(ONNXModel, factory) 17 | assert opset_lhs, "Should not be empty... Something must go wrong." 18 | opset_rhs = auto_opconfig(ONNXModel, factory) 19 | assert opset_lhs == opset_rhs 20 | 21 | # Assert types 22 | assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType) 23 | 24 | # Assert Dictionary Type Equality 25 | assert type(opset_lhs) == type(opset_rhs) 26 | assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"]) 27 | assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type( 28 | opset_rhs["core.ReLU"].in_dtypes[0][0] 29 | ) 30 | 31 | 32 | def test_synthesized_onnx_model(tmp_path): 33 | d = tmp_path / "test_tvm_onnx" 34 | d.mkdir() 35 | 36 | ONNXModel = Model.init("onnx") 37 | factory = BackendFactory.init( 38 | "tvm", 39 | target="cuda" if tvm.cuda(0).exist else "cpu", 40 | optmax=False, 41 | ) 42 | 43 | gen = model_gen( 44 | opset=auto_opset(ONNXModel, factory), 45 | seed=23132, 46 | max_nodes=1, 47 | ) # One op should not be easily wrong... I guess. 48 | 49 | model = ONNXModel.from_gir(gen.make_concrete()) 50 | 51 | assert model.with_torch 52 | 53 | model.refine_weights() # either random generated or gradient-based. 54 | oracle = model.make_oracle() 55 | 56 | testcase = TestCase(model, oracle) 57 | testcase.dump(root_folder=d) 58 | 59 | assert factory.verify_testcase(testcase) is None 60 | -------------------------------------------------------------------------------- /exp_helper/1_sumarize_buglist_raw.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pandas as pd 3 | 4 | RAW_BUGLIST_PATH = "/root/raw_data/raw_data/tab3_raw_data/polyjuice_buglist.csv" 5 | OUTPUT_DIR_PATH = "/root/raw_data/raw_data_results" 6 | 7 | 8 | def main(): 9 | with open(RAW_BUGLIST_PATH, "r") as f: 10 | buglist = f.readlines() 11 | 12 | res_dic = { 13 | "Torch Inductor": {"reported": 0, "confirmed": 0, "fixed": 0}, 14 | "TensorRT": {"reported": 0, "confirmed": 0, "fixed": 0}, 15 | "ONNXRuntime": {"reported": 0, "confirmed": 0, "fixed": 0}, 16 | "TVM": {"reported": 0, "confirmed": 0, "fixed": 0}, 17 | "TF XLA": {"reported": 0, "confirmed": 0, "fixed": 0}, 18 | "Hidet": {"reported": 0, "confirmed": 0, "fixed": 0}, 19 | "EinNet": {"reported": 0, "confirmed": 0, "fixed": 0}, 20 | } 21 | 22 | for bug_report in buglist[1:]: 23 | col = bug_report.split(",") 24 | assert len(col) == 3 25 | compiler_name = col[0].split("#")[0] 26 | status = col[2].strip() 27 | if status == "reported": 28 | res_dic[compiler_name]["reported"] += 1 29 | elif status == "confirmed": 30 | res_dic[compiler_name]["reported"] += 1 31 | res_dic[compiler_name]["confirmed"] += 1 32 | elif status == "fixed": 33 | res_dic[compiler_name]["reported"] += 1 34 | res_dic[compiler_name]["confirmed"] += 1 35 | res_dic[compiler_name]["fixed"] += 1 36 | else: 37 | raise AssertionError(f"Unknown status: {status}") 38 | 39 | res_list = [] 40 | for compiler_name, status_dic in res_dic.items(): 41 | res_list.append([compiler_name, status_dic["reported"], status_dic["confirmed"], status_dic["fixed"]]) 42 | 43 | output_path = os.path.join(OUTPUT_DIR_PATH, "tab3.csv") 44 | 45 | pd.DataFrame(data=res_list, columns=["Compiler", "Reported", "Confirmed", "Fixed"]).to_csv( 46 | output_path, 47 | index=False) 48 | 49 | 50 | if __name__ == '__main__': 51 | main() 52 | -------------------------------------------------------------------------------- /tests/tensorflow/test_xla_backend.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import tensorflow as tf 3 | 4 | from nnsmith.abstract.dtype import DType 5 | from nnsmith.backends import BackendFactory 6 | from nnsmith.graph_gen import model_gen 7 | from nnsmith.materialize import Model, TestCase 8 | from nnsmith.narrow_spec import auto_opconfig, auto_opset 9 | 10 | TestCase.__test__ = False # supress PyTest warning 11 | 12 | 13 | def test_narrow_spec_cache_make_and_reload(): 14 | factory = BackendFactory.init("xla", target="cpu", optmax=True) 15 | ModelType = Model.init("tensorflow") 16 | opset_lhs = auto_opconfig(ModelType, factory) 17 | assert opset_lhs, "Should not be empty... Something must go wrong." 18 | opset_rhs = auto_opconfig(ModelType, factory) 19 | assert opset_lhs == opset_rhs 20 | 21 | # Assert types 22 | assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType) 23 | 24 | # Assert Dictionary Type Equality 25 | assert type(opset_lhs) == type(opset_rhs) 26 | assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"]) 27 | assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type( 28 | opset_rhs["core.ReLU"].in_dtypes[0][0] 29 | ) 30 | 31 | 32 | def test_synthesized_model(tmp_path): 33 | d = tmp_path / "test_xla" 34 | d.mkdir() 35 | 36 | targets = ["cpu"] 37 | if tf.config.list_logical_devices("GPU"): 38 | targets.append("cuda") 39 | 40 | for target in targets: 41 | factory = BackendFactory.init("xla", target=target, optmax=False) 42 | 43 | ModelType = Model.init("tensorflow", backend_target=target) 44 | 45 | gen = model_gen( 46 | opset=auto_opset(ModelType, factory), 47 | seed=23132, 48 | max_nodes=4, 49 | ) # One op should not be easily wrong... I guess. 50 | 51 | model = ModelType.from_gir(gen.make_concrete()) 52 | 53 | oracle = model.make_oracle() 54 | 55 | testcase = TestCase(model, oracle) 56 | assert factory.verify_testcase(testcase) is None 57 | testcase.dump(root_folder=d) 58 | -------------------------------------------------------------------------------- /exp_helper/process_diff_effect.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | from scipy import stats 3 | import sys 4 | import os 5 | 6 | pd.set_option('display.max_columns', None) 7 | pd.set_option('display.width', 100) 8 | 9 | 10 | def process_diff_effect(path): 11 | df = pd.read_csv(path) 12 | 13 | naive_edit_dis_mean = df["edit_res_1"].mean() 14 | polyjuice_edit_dis_mean = df["edit_res_2"].mean() 15 | 16 | naive_lcs_dis_mean = df["lcs_res_1"].mean() 17 | polyjuice_lcs_dis_mean = df["lcs_res_2"].mean() 18 | 19 | df_tmp = df[df['edit_res_1'] != 0] 20 | df2 = df_tmp[df_tmp['lcs_res_1'] != 0] 21 | 22 | edit_dis_impr = ((df2['edit_res_2'] - df2['edit_res_1']) / df2['edit_res_1']) 23 | lcs_dis_impr = ((df2['lcs_res_2'] - df2['lcs_res_1']) / df2['lcs_res_1']) 24 | 25 | edit_dis_avg_impr = edit_dis_impr.mean() 26 | lcs_dis_avg_impr = lcs_dis_impr.mean() 27 | 28 | edit_dis_median_impr = edit_dis_impr.median() 29 | lcs_dis_median_impr = lcs_dis_impr.median() 30 | 31 | edit_dis_impr_std = edit_dis_impr.std() 32 | lcs_dis_impr_std = lcs_dis_impr.std() 33 | 34 | edit_dis_impr_stderr = edit_dis_impr.sem() 35 | lcs_dis_impr_stderr = lcs_dis_impr.sem() 36 | 37 | res_df = pd.DataFrame(data=[ 38 | ["{0:.2f}%".format(lcs_dis_avg_impr * 100), lcs_dis_impr_stderr, "{0:.2f}%".format(edit_dis_avg_impr * 100), 39 | edit_dis_impr_stderr]], 40 | columns=["lcs_dis_avg_impr", "lcs_dis_impr_stderr", "edit_dis_avg_impr", 41 | "edit_dis_impr_stderr", ]) 42 | print(res_df) 43 | 44 | output_dir_path = os.environ.get("OUTPUT_DIR") 45 | if output_dir_path is not None: 46 | csv_output_path = os.path.join(output_dir_path, "table4.csv") 47 | res_df.to_csv(csv_output_path, index=False) 48 | 49 | # print({ 50 | # "edit_dis_avg_impr": edit_dis_avg_impr, "lcs_dis_avg_impr": lcs_dis_avg_impr, 51 | # "edit_dis_impr_stderr": edit_dis_impr_stderr, "lcs_dis_impr_stderr": lcs_dis_impr_stderr}) 52 | 53 | 54 | if __name__ == '__main__': 55 | path = sys.argv[1] 56 | process_diff_effect(path) 57 | -------------------------------------------------------------------------------- /tests/torch/test_pt2_backend.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nnsmith.abstract.dtype import DType 5 | from nnsmith.backends import BackendFactory 6 | from nnsmith.graph_gen import model_gen 7 | from nnsmith.materialize import Model, TestCase 8 | from nnsmith.narrow_spec import auto_opconfig, auto_opset 9 | 10 | TestCase.__test__ = False # supress PyTest warning 11 | 12 | 13 | def test_narrow_spec_cache_make_and_reload(): 14 | factory = BackendFactory.init("pt2", target="cpu", optmax=True) 15 | ModelType = Model.init("torch") 16 | opset_lhs = auto_opconfig(ModelType, factory) 17 | assert opset_lhs, "Should not be empty... Something must go wrong." 18 | opset_rhs = auto_opconfig(ModelType, factory) 19 | assert opset_lhs == opset_rhs 20 | 21 | # Assert types 22 | assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType) 23 | 24 | # Assert Dictionary Type Equality 25 | assert type(opset_lhs) == type(opset_rhs) 26 | assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"]) 27 | assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type( 28 | opset_rhs["core.ReLU"].in_dtypes[0][0] 29 | ) 30 | 31 | 32 | def test_synthesized_model(tmp_path): 33 | d = tmp_path / "test_pt2" 34 | d.mkdir() 35 | 36 | targets = ["cpu"] 37 | if torch.cuda.is_available(): 38 | targets.append("cuda") 39 | 40 | for target in targets: 41 | for grad in [True, False]: 42 | factory = BackendFactory.init("pt2", target=target, optmax=False) 43 | 44 | ModelType = Model.init("torch", backend_target=target) 45 | 46 | gen = model_gen( 47 | opset=auto_opset(ModelType, factory, grad=grad), 48 | seed=23132, 49 | max_nodes=1, 50 | ) # One op should not be easily wrong... I guess. 51 | 52 | model = ModelType.from_gir(gen.make_concrete()) 53 | 54 | model.set_grad_check(grad) 55 | oracle = model.make_oracle() 56 | 57 | testcase = TestCase(model, oracle) 58 | assert factory.verify_testcase(testcase) is None 59 | testcase.dump(root_folder=d) 60 | -------------------------------------------------------------------------------- /tests/tensorrt/test_trt_backend.py: -------------------------------------------------------------------------------- 1 | import GPUtil 2 | import pytest 3 | 4 | has_gpu = len(GPUtil.getGPUs()) > 0 5 | 6 | from nnsmith.abstract.dtype import DType 7 | from nnsmith.backends import BackendFactory 8 | from nnsmith.graph_gen import model_gen 9 | from nnsmith.materialize import Model, TestCase 10 | from nnsmith.narrow_spec import auto_opconfig, auto_opset 11 | 12 | TestCase.__test__ = False # supress PyTest warning 13 | 14 | 15 | @pytest.mark.skipif(not has_gpu, reason="Skipping TensorRT testing due to no GPU found") 16 | def test_narrow_spec_cache_make_and_reload(): 17 | factory = BackendFactory.init("tensorrt", target="cuda", optmax=True) 18 | ONNXModel = Model.init("onnx") 19 | opset_lhs = auto_opconfig(ONNXModel, factory) 20 | assert opset_lhs, "Should not be empty... Something must go wrong." 21 | opset_rhs = auto_opconfig(ONNXModel, factory) 22 | assert opset_lhs == opset_rhs 23 | 24 | # Assert types 25 | assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType) 26 | 27 | # Assert Dictionary Type Equality 28 | assert type(opset_lhs) == type(opset_rhs) 29 | assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"]) 30 | assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type( 31 | opset_rhs["core.ReLU"].in_dtypes[0][0] 32 | ) 33 | 34 | 35 | @pytest.mark.skipif(not has_gpu, reason="Skipping TensorRT testing due to no GPU found") 36 | def test_synthesized_onnx_model(tmp_path): 37 | d = tmp_path / "test_trt_onnx" 38 | d.mkdir() 39 | 40 | ONNXModel = Model.init("onnx") 41 | factory = BackendFactory.init("tensorrt", target="cuda", optmax=True) 42 | 43 | gen = model_gen( 44 | opset=auto_opset(ONNXModel, factory), 45 | seed=23132, 46 | max_nodes=1, 47 | ) # One op should not be easily wrong... I guess. 48 | 49 | model = ONNXModel.from_gir(gen.make_concrete()) 50 | 51 | assert model.with_torch 52 | 53 | model.refine_weights() # either random generated or gradient-based. 54 | oracle = model.make_oracle() 55 | 56 | testcase = TestCase(model, oracle) 57 | testcase.dump(root_folder=d) 58 | 59 | assert factory.verify_testcase(testcase) is None 60 | -------------------------------------------------------------------------------- /tests/torch/test_torchjit_backend.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nnsmith.abstract.dtype import DType 5 | from nnsmith.backends import BackendFactory 6 | from nnsmith.graph_gen import model_gen 7 | from nnsmith.materialize import Model, TestCase 8 | from nnsmith.narrow_spec import auto_opconfig, auto_opset 9 | 10 | TestCase.__test__ = False # supress PyTest warning 11 | 12 | 13 | def test_narrow_spec_cache_make_and_reload(): 14 | factory = BackendFactory.init("torchjit", target="cpu", optmax=True) 15 | ModelType = Model.init("torch") 16 | opset_lhs = auto_opconfig(ModelType, factory) 17 | assert opset_lhs, "Should not be empty... Something must go wrong." 18 | opset_rhs = auto_opconfig(ModelType, factory) 19 | assert opset_lhs == opset_rhs 20 | 21 | # Assert types 22 | assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType) 23 | 24 | # Assert Dictionary Type Equality 25 | assert type(opset_lhs) == type(opset_rhs) 26 | assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"]) 27 | assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type( 28 | opset_rhs["core.ReLU"].in_dtypes[0][0] 29 | ) 30 | 31 | 32 | def test_synthesized_model(tmp_path): 33 | d = tmp_path / "test_torchjit" 34 | d.mkdir() 35 | 36 | targets = ["cpu"] 37 | if torch.cuda.is_available(): 38 | targets.append("cuda") 39 | 40 | for target in targets: 41 | for grad in [True, False]: 42 | factory = BackendFactory.init("torchjit", target=target, optmax=False) 43 | 44 | ModelType = Model.init("torch", backend_target=target) 45 | 46 | gen = model_gen( 47 | opset=auto_opset(ModelType, factory, grad=grad), 48 | seed=23132, 49 | max_nodes=1, 50 | ) # One op should not be easily wrong... I guess. 51 | 52 | model = ModelType.from_gir(gen.make_concrete()) 53 | 54 | model.set_grad_check(grad) 55 | oracle = model.make_oracle() 56 | 57 | testcase = TestCase(model, oracle) 58 | assert factory.verify_testcase(testcase) is None 59 | testcase.dump(root_folder=d) 60 | -------------------------------------------------------------------------------- /exp_helper/3_run_eq_effectiveness.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import shutil 3 | import os 4 | import copy 5 | 6 | POLYJUICE_ROOT_PATH = "/root/polyjuice/" 7 | TMP_FUZZ_ROOT = "/tmp/fuzz_report" 8 | 9 | OUTPUT_DIR = "/root/3_eq_effectiveness/" 10 | 11 | NOISE_REDIRECTION = subprocess.STDOUT 12 | 13 | 14 | def main(): 15 | if os.path.exists(OUTPUT_DIR): 16 | shutil.rmtree(OUTPUT_DIR) 17 | 18 | os.makedirs(OUTPUT_DIR, exist_ok=True) 19 | subprocess.run(["rm", "-rf", TMP_FUZZ_ROOT + "*"]) 20 | 21 | new_env = copy.deepcopy(os.environ) 22 | new_env["OUTPUT_DIR"] = OUTPUT_DIR 23 | new_env["SAVE_EQ_IS_WRONG"] = "false" 24 | new_env["EQ_HELPER_PATH"] = "/root/polyjuice/libdl_compiler_fuzzer_helper.so" 25 | new_env["TVM_LOG_DEBUG"] = "DEFAULT=0" 26 | new_env["TVM_HOME"] = "/root/compilers/apache-tvm-src" 27 | new_env["PYTHONPATH"] = "/root/compilers/apache-tvm-src/python:/root/polyjuice" 28 | new_env["LD_LIBRARY_PATH"] = "/root/compilers/apache-tvm-src/new_build" 29 | 30 | ablation_study_script_path = POLYJUICE_ROOT_PATH + "nnsmith/cli/ablation_exp_run.py" 31 | p = subprocess.run( 32 | ["python3", ablation_study_script_path, "mgen.rank_choices=[3,3]", "model.type=onnx", "backend.type=tvm", 33 | "backend.target=cpu", f"fuzz.root={TMP_FUZZ_ROOT}", "debug.viz=true", "mgen.max_nodes=10"], 34 | # stdout=NOISE_REDIRECTION, stderr=NOISE_REDIRECTION, 35 | env=new_env, cwd=POLYJUICE_ROOT_PATH) 36 | if p.returncode != 0: 37 | print(f"ablation study failed. return code: {p.returncode}") 38 | else: 39 | print("ablation study runs finished successfully") 40 | print("processing results...") 41 | 42 | process_diff_script_path = os.path.join(POLYJUICE_ROOT_PATH, "exp_helper/process_diff_effect.py") 43 | ablation_exp_csv_path = os.path.join(OUTPUT_DIR, "ablation_exp.csv") 44 | 45 | p = subprocess.run( 46 | ["python3", process_diff_script_path, ablation_exp_csv_path], 47 | # stdout=NOISE_REDIRECTION, stderr=NOISE_REDIRECTION, 48 | env=new_env, cwd=POLYJUICE_ROOT_PATH) 49 | if p.returncode != 0: 50 | print(f"processing results failed. return code: {p.returncode}") 51 | else: 52 | print("processing results finished successfully") 53 | 54 | print(f"the result should be in {os.path.join(OUTPUT_DIR, 'table4.csv')}") 55 | 56 | 57 | if __name__ == '__main__': 58 | main() 59 | -------------------------------------------------------------------------------- /tests/core/test_parse_name_kwargs.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from nnsmith.backends.factory import parse_name_kwargs 4 | 5 | 6 | def test_single_invalid(): 7 | with pytest.raises(ValueError, match="Invalid backend"): 8 | parse_name_kwargs("") 9 | 10 | with pytest.raises(ValueError, match="Invalid backend"): 11 | parse_name_kwargs("+") 12 | 13 | with pytest.raises(ValueError, match="Invalid backend"): 14 | parse_name_kwargs("@something") 15 | 16 | with pytest.raises(ValueError, match="Invalid backend"): 17 | parse_name_kwargs("something@") 18 | 19 | with pytest.raises(ValueError, match="Invalid backend"): 20 | parse_name_kwargs("pt2@something") 21 | 22 | 23 | def test_single_valid(): 24 | assert parse_name_kwargs("pt2") == ("pt2", {}) 25 | assert parse_name_kwargs("pt2 ") == ("pt2", {}) 26 | assert parse_name_kwargs("pt2 ") == ("pt2", {}) 27 | assert parse_name_kwargs(" pt2") == ("pt2", {}) 28 | assert parse_name_kwargs(" pt2") == ("pt2", {}) 29 | assert parse_name_kwargs(" pt2 ") == ("pt2", {}) 30 | 31 | 32 | def test_kwargs_invalid(): 33 | with pytest.raises(ValueError, match="Invalid backend"): 34 | parse_name_kwargs("pt2 foo") 35 | 36 | with pytest.raises(ValueError, match="Invalid backend"): 37 | parse_name_kwargs("pt2 foo@") 38 | 39 | with pytest.raises(ValueError, match="Invalid backend"): 40 | parse_name_kwargs("pt2 @bar") 41 | 42 | with pytest.raises(ValueError, match="Invalid backend"): 43 | parse_name_kwargs("pt2 foo@bar baz") 44 | 45 | with pytest.raises(ValueError, match="Invalid backend"): 46 | parse_name_kwargs("pt2 foo@ qux") 47 | 48 | with pytest.raises(ValueError, match="Invalid backend"): 49 | parse_name_kwargs("pt2 foo@@qux") 50 | 51 | 52 | def test_kwargs_valid(): 53 | assert parse_name_kwargs("pt2 foo@bar") == ("pt2", {"foo": "bar"}) 54 | 55 | assert parse_name_kwargs("pt2 foo@bar baz@qux") == ( 56 | "pt2", 57 | {"foo": "bar", "baz": "qux"}, 58 | ) 59 | 60 | assert parse_name_kwargs("pt2 foo@bar baz@qux quux@quuz") == ( 61 | "pt2", 62 | {"foo": "bar", "baz": "qux", "quux": "quuz"}, 63 | ) 64 | 65 | # add a few random space 66 | assert parse_name_kwargs("pt2 foo@bar baz@qux quux@quuz") == ( 67 | "pt2", 68 | {"foo": "bar", "baz": "qux", "quux": "quuz"}, 69 | ) 70 | -------------------------------------------------------------------------------- /nnsmith/error.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from abc import ABC, abstractmethod 3 | 4 | 5 | class InternalError(Exception): 6 | """Fatal unexpected internal errors in NNSmith that should shut down the program immediately.""" 7 | 8 | pass 9 | 10 | 11 | class ConstraintError(Exception): 12 | """Expected possible constarint unsat error used in shape transfer function.""" 13 | 14 | pass 15 | 16 | 17 | class BaseChecker(ABC): 18 | @classmethod 19 | @abstractmethod 20 | def handler(cls, msg): 21 | pass 22 | 23 | @classmethod 24 | def eq(cls, lhs, rhs, msg=""): 25 | if lhs != rhs: 26 | cls.handler(f"Failed asertion :: {msg} | {lhs} != {rhs}") 27 | 28 | @classmethod 29 | def gt(cls, lhs, rhs, msg=""): 30 | if lhs <= rhs: 31 | cls.handler(f"Failed asertion :: {msg} | {lhs} <= {rhs}") 32 | 33 | @classmethod 34 | def ge(cls, lhs, rhs, msg=""): 35 | if lhs < rhs: 36 | cls.handler(f"Failed asertion :: {msg} | {lhs} < {rhs}") 37 | 38 | @classmethod 39 | def lt(cls, lhs, rhs, msg=""): 40 | if lhs >= rhs: 41 | cls.handler(f"Failed asertion :: {msg} | {lhs} >= {rhs}") 42 | 43 | @classmethod 44 | def le(cls, lhs, rhs, msg=""): 45 | if lhs > rhs: 46 | cls.handler(f"Failed asertion :: {msg} | {lhs} > {rhs}") 47 | 48 | @classmethod 49 | def none(cls, obj, msg=""): 50 | if obj is not None: 51 | cls.handler(f"Failed asertion :: {msg} | expr is not None") 52 | 53 | @classmethod 54 | def not_none(cls, obj, msg=""): 55 | if obj is None: 56 | cls.handler(f"Failed asertion :: {msg} | expr is None") 57 | 58 | @classmethod 59 | def true(cls, cond, msg=""): 60 | if not cond: 61 | cls.handler(f"Failed asertion :: {msg} | condition is not True") 62 | 63 | @classmethod 64 | def false(cls, cond, msg=""): 65 | if cond: 66 | cls.handler(f"Failed asertion :: {msg} | condition is not False") 67 | 68 | 69 | class SanityCheck(BaseChecker): 70 | @classmethod 71 | def handler(cls, msg): 72 | logging.critical(msg) 73 | raise InternalError( 74 | msg + " | Reporting bugs @ https://github.com/ise-uiuc/nnsmith/issues" 75 | ) 76 | 77 | 78 | class ConstraintCheck(BaseChecker): 79 | @classmethod 80 | def handler(cls, msg): 81 | raise ConstraintError(msg) 82 | -------------------------------------------------------------------------------- /nnsmith/backends/infinitensor.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | import onnx 4 | from pyinfinitensor.onnx import OnnxStub 5 | from multipledispatch import dispatch 6 | 7 | from nnsmith.backends import BackendFactory 8 | from nnsmith.backends.factory import BackendCallable 9 | from nnsmith.macro import NNSMITH_ORT_INTRA_OP_THREAD 10 | from nnsmith.materialize.onnx import ONNXModel 11 | from nnsmith.abstract.dtype import DType 12 | 13 | class IT(BackendFactory): 14 | def __init__(self, target, optmax, **kwargs): 15 | """opt_level ranges from 0 to 3, stands for ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_EXTENDED and ORT_ENABLE_ALL. 16 | See https://onnxruntime.ai/docs/performance/graph-optimizations.html for detail 17 | """ 18 | super().__init__(target, optmax, **kwargs) 19 | 20 | if target == "cuda": 21 | from pyinfinitensor.onnx import backend 22 | self.runtime = backend.cuda_runtime() 23 | elif target == "cpu": 24 | from pyinfinitensor import backend 25 | self.runtime = backend.cpu_runtime() 26 | else: 27 | raise ValueError( 28 | f"Unknown target `{target}`. Only `cpu` and `cuda` are supported." 29 | ) 30 | 31 | @property 32 | def system_name(self) -> str: 33 | return "infinitetensor" 34 | 35 | @property 36 | def import_libs(self) -> List[str]: 37 | return ["from pyinfinitensor.onnx import OnnxStub", "from pyinfinitensor import backend"] 38 | 39 | @classmethod 40 | def skip_dtypes(cls) -> List[DType]: 41 | # TRT will truncate f64 -> f32 and i64 -> i32 42 | return [DType.float64, DType.float16, DType.uint64, DType.uint16, DType.uint8] 43 | 44 | @dispatch(ONNXModel) 45 | def make_backend( 46 | self, 47 | model: ONNXModel, 48 | ) -> BackendCallable: 49 | onnx_model = model.native_model 50 | stub = OnnxStub(onnx_model, self.runtime) 51 | 52 | out_names = list(model.output_like.keys()) 53 | 54 | def closure(inputs): 55 | 56 | for name, tensor in stub.inputs.items(): 57 | tmp_input = inputs[name] 58 | tensor.copyin_numpy(tmp_input) 59 | 60 | stub.run() 61 | res = [v.copyout_numpy() for v in stub.outputs.values()] 62 | return {n: r for n, r in zip(out_names, res)} 63 | 64 | return closure 65 | -------------------------------------------------------------------------------- /experiments/legacy/plot_cov.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import matplotlib.pyplot as plt 3 | import pandas 4 | import os 5 | 6 | 7 | class Ploter: 8 | def __init__(self, cov_lim=None) -> None: 9 | self.legends = [] # type: ignore 10 | # cov / time, cov / iteration, iteration / time 11 | fig, axs = plt.subplots(1, 3, constrained_layout=True, figsize=(16, 5)) 12 | self.fig = fig 13 | self.axs = axs 14 | self.cov_lim = cov_lim 15 | 16 | def add(self, folder, name=None): 17 | path = os.path.join(folder, "cov_by_time.csv") 18 | df = pandas.read_csv(path, usecols=[0, 1], header=None).to_numpy() 19 | 20 | self.axs[0].plot(df[:, 0], df[:, 1]) # cov / time 21 | self.axs[1].plot(range(len(df[:, 0])), df[:, 1]) # cov / iteration 22 | self.axs[2].plot(df[:, 0], range(len(df[:, 1]))) # iter / time 23 | 24 | if name: 25 | self.legends.append(name) 26 | else: 27 | assert not self.legends 28 | 29 | def plot(self, save="cov"): 30 | for axs in self.axs: 31 | axs.legend(self.legends) 32 | # plt.legend(self.legends) 33 | 34 | if self.cov_lim is not None: 35 | self.axs[0].set_ylim(bottom=self.cov_lim) 36 | self.axs[1].set_ylim(bottom=self.cov_lim) 37 | 38 | self.axs[0].set(xlabel="Time / Second", ylabel="# Coverage") 39 | self.axs[0].set_title("Coverage $\\bf{Time}$ Efficiency") 40 | 41 | self.axs[1].set(ylabel="# Coverage", xlabel="# Iteration") 42 | self.axs[1].set_title("Coverage $\\bf{Iteration}$ Efficiency") 43 | 44 | self.axs[2].set(xlabel="Time / Second", ylabel="# Iteration") 45 | self.axs[2].set_title("Iteration Speed") 46 | 47 | plt.savefig(save + ".pdf") 48 | plt.savefig(save + ".png") 49 | 50 | 51 | if "__main__" == __name__: 52 | parser = argparse.ArgumentParser() 53 | parser.add_argument( 54 | "-f", "--folders", type=str, nargs="+", help="bug report folder" 55 | ) 56 | parser.add_argument( 57 | "-cl", "--cov_lim", type=int, default=None, help="coverage starting lim" 58 | ) 59 | parser.add_argument( 60 | "--tvmfuzz", type=str, nargs="?", help="TVMFuzz coverage by time file" 61 | ) 62 | args = parser.parse_args() 63 | 64 | ploter = Ploter(cov_lim=args.cov_lim) 65 | 66 | for f in args.folders: 67 | ploter.add(f, f) 68 | ploter.plot("cov") 69 | -------------------------------------------------------------------------------- /nnsmith/backends/tensorrt.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import List 3 | 4 | import numpy as np 5 | import onnx 6 | import pycuda.driver as cuda 7 | import tensorrt as trt 8 | from multipledispatch import dispatch 9 | from pycuda.driver import DeviceAllocation 10 | 11 | from nnsmith.abstract.arith import * 12 | from nnsmith.abstract.dtype import DType 13 | from nnsmith.abstract.extension import patch_requires 14 | from nnsmith.abstract.op import AbsOpBase 15 | from nnsmith.abstract.tensor import AbsTensor 16 | from nnsmith.backends import BackendFactory 17 | from nnsmith.materialize.onnx import ONNXModel 18 | from polygraphy.backend.trt import EngineFromNetwork, TrtRunner, NetworkFromOnnxBytes 19 | from polygraphy.logger.logger import G_LOGGER 20 | G_LOGGER.module_severity = G_LOGGER.ERROR 21 | 22 | 23 | @dataclass 24 | class HostDeviceMem: 25 | host: np.ndarray 26 | device: DeviceAllocation 27 | 28 | 29 | class TRT(BackendFactory): 30 | def __init__(self, target="cuda", optmax=True, **kwargs): 31 | super().__init__(target, optmax, **kwargs) 32 | 33 | if target != "cuda": 34 | raise ValueError("TensorRT backend only supports GPU!") 35 | 36 | if optmax is False: 37 | # TODO(@ganler): support non-optimized TensorRT by using performing 38 | # inference over a model that marks all nodes as outputs. 39 | raise ValueError("There is not O0 mode for TensorRT so far.") 40 | 41 | @property 42 | def system_name(self) -> str: 43 | return "tensorrt" 44 | 45 | @property 46 | def version(self) -> str: 47 | return trt.__version__ 48 | 49 | @dispatch(ONNXModel) 50 | def make_backend(self, model: ONNXModel): 51 | engine = EngineFromNetwork(NetworkFromOnnxBytes(model.native_model.SerializeToString())) 52 | 53 | def closure(inputs): 54 | with TrtRunner(engine) as runner: 55 | return runner.infer(inputs) 56 | 57 | return closure 58 | 59 | @property 60 | def import_libs(self) -> List[str]: 61 | return ["import tensorrt as trt"] 62 | 63 | @classmethod 64 | def skip_dtypes(cls) -> List[DType]: 65 | # TRT will truncate f64 -> f32 and i64 -> i32 66 | return [DType.float64, DType.int64] 67 | 68 | 69 | @patch_requires(TRT.system_name, "core.Pool2d") 70 | def RulePool2d(self: AbsOpBase, _: List[AbsTensor]) -> List[Union[z3.BoolRef, bool]]: 71 | return [nnsmith_lt(nnsmith_mul(self.kh, self.kw), 10000)] 72 | -------------------------------------------------------------------------------- /tests/torch/test_biconvert.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nnsmith.materialize.torch.parse import parse 5 | from nnsmith.materialize.torch.symbolnet import FxTracing, SymbolNet 6 | 7 | 8 | def test_biconvert(): 9 | class MyModel(torch.nn.Module): 10 | def __init__( 11 | self, 12 | ): 13 | super().__init__() 14 | self.linear = torch.nn.Linear(3, 4) 15 | 16 | def forward(self, i0, i1): 17 | v0 = i0 + 3.14 + i1[0, 0] 18 | v1 = self.linear(v0) 19 | v1_0, v1_1 = torch.split(v1, [1, 3], dim=-1) 20 | v2 = torch.mul(input=v1_0, other=v1_1) 21 | v3 = torch.cat([v2, v2], dim=-1) 22 | v4 = v3.flatten() 23 | return v4 24 | 25 | model = MyModel() 26 | i0 = torch.rand(2, 3) 27 | i1 = torch.rand(1, 2) 28 | ir = parse(model, i0, i1) 29 | assert ( 30 | ir.pretty().strip() 31 | == """\ 32 | v0_0 = Input() # inst id: 0 33 | v1_0 = core.ConcreteOp(v0_0) # inst id: 1 34 | v2_0 = Input() # inst id: 2 35 | v3_0 = core.ConcreteOp(v2_0) # inst id: 3 36 | v4_0 = core.ConcreteOp(v3_0, v1_0) # inst id: 4 37 | v5_0 = core.ConcreteOp(v4_0) # inst id: 5 38 | v6_0, v6_1 = core.ConcreteOp(v5_0) # inst id: 6 39 | v7_0 = core.ConcreteOp(v6_0, v6_1) # inst id: 7 40 | v8_0 = core.ConcreteOp(v7_0, v7_0) # inst id: 8 41 | v9_0 = core.ConcreteOp(v8_0) # inst id: 9""" 42 | ) 43 | 44 | ir.remove_unused(ir.insts[-1]) # mutate: remove the last flatten op. 45 | 46 | net = SymbolNet(ir) 47 | with FxTracing(): 48 | traced = torch.fx.symbolic_trace(net) 49 | assert ( 50 | traced.code.strip() 51 | == R"""def forward(self, *args): 52 | _args = args 53 | getitem = _args[0] 54 | getitem_1 = _args[1]; _args = None 55 | getitem_2 = getitem[(0, 0)]; getitem = None 56 | add = getitem_1 + 3.14; getitem_1 = None 57 | add_1 = add + getitem_2; add = getitem_2 = None 58 | m5 = self.m5(add_1); add_1 = None 59 | split = torch.functional.split(m5, [1, 3], dim = -1); m5 = None 60 | getitem_3 = split[0] 61 | getitem_4 = split[0] 62 | getitem_5 = split[1] 63 | getitem_6 = split[1]; split = None 64 | mul = torch.mul(input = getitem_3, other = getitem_5); getitem_3 = getitem_5 = None 65 | cat = torch.cat([mul, mul], dim = -1); mul = None 66 | return (cat,)""" 67 | ) 68 | -------------------------------------------------------------------------------- /nnsmith/backends/onnxruntime.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | import onnx 4 | import onnxruntime as ort 5 | from multipledispatch import dispatch 6 | 7 | from nnsmith.backends import BackendFactory 8 | from nnsmith.backends.factory import BackendCallable 9 | from nnsmith.macro import NNSMITH_ORT_INTRA_OP_THREAD 10 | from nnsmith.materialize.onnx import ONNXModel 11 | 12 | OPT_LEVELS = [ 13 | ort.GraphOptimizationLevel.ORT_DISABLE_ALL, 14 | ort.GraphOptimizationLevel.ORT_ENABLE_BASIC, 15 | ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED, 16 | ort.GraphOptimizationLevel.ORT_ENABLE_ALL, 17 | ] 18 | 19 | 20 | class ORT(BackendFactory): 21 | def __init__(self, target, optmax, **kwargs): 22 | """opt_level ranges from 0 to 3, stands for ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_EXTENDED and ORT_ENABLE_ALL. 23 | See https://onnxruntime.ai/docs/performance/graph-optimizations.html for detail 24 | """ 25 | super().__init__(target, optmax, **kwargs) 26 | self.opt_level = OPT_LEVELS[-1 if optmax else 0] 27 | 28 | if target == "cuda": 29 | self.providers = [ 30 | "CUDAExecutionProvider", 31 | "CPUExecutionProvider", 32 | ] # ordered by precedence 33 | elif target == "cpu": 34 | self.providers = ["CPUExecutionProvider"] 35 | else: 36 | raise ValueError( 37 | f"Unknown target `{target}`. Only `cpu` and `cuda` are supported." 38 | ) 39 | 40 | @property 41 | def system_name(self) -> str: 42 | return "onnxruntime" 43 | 44 | @property 45 | def version(self) -> str: 46 | return ort.__version__ 47 | 48 | @property 49 | def import_libs(self) -> List[str]: 50 | return ["import onnxruntime as ort"] 51 | 52 | @dispatch(ONNXModel) 53 | def make_backend( 54 | self, 55 | model: ONNXModel, 56 | ) -> BackendCallable: 57 | sess_options = ort.SessionOptions() 58 | sess_options.graph_optimization_level = self.opt_level 59 | # https://github.com/microsoft/onnxruntime/issues/8313 60 | sess_options.intra_op_num_threads = NNSMITH_ORT_INTRA_OP_THREAD 61 | 62 | sess = ort.InferenceSession( 63 | onnx._serialize(model.native_model), 64 | providers=self.providers, 65 | sess_options=sess_options, 66 | ) 67 | out_names = list(model.output_like.keys()) 68 | 69 | def closure(inputs): 70 | res = sess.run(out_names, inputs) 71 | return {n: r for n, r in zip(out_names, res)} 72 | 73 | return closure 74 | -------------------------------------------------------------------------------- /experiments/legacy/compare_tzer.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import argparse 3 | import os 4 | 5 | import numpy as np 6 | 7 | from tvm.contrib import coverage 8 | 9 | 10 | def clipper(cov): 11 | # Okay, fine, due my previous stupid implementation: 12 | # https://github.com/ganler/tvm/blob/7e298af66ced5216846a6eb9f9bc01b677cf05d5/python/tvm/contrib/coverage.py#L19 13 | # the byte array is 8 times larger than the coverage (was a bit array) as I forgot to "/8" 14 | # but note that this won't affect the results but just waste 7x space... 15 | # So what I'm gonna do here is to clip it to the required size. 16 | cov_length = coverage.get_total() 17 | required_bytes = (cov_length + 7) // 8 18 | 19 | return np.unpackbits(cov[:required_bytes])[:cov_length] # you got it. 20 | 21 | 22 | if __name__ == "__main__": 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument( 25 | "--tzer", type=str, required=True, help="Folder of tzer reports." 26 | ) 27 | parser.add_argument( 28 | "--nnsmith", 29 | type=str, 30 | required=True, 31 | help="Folder of nnsmith evaluated in memcov.", 32 | ) 33 | args = parser.parse_args() 34 | 35 | with open(os.path.join(args.tzer, "cov.pkl"), "rb") as fp: 36 | tzer_cov = pickle.load(fp) 37 | 38 | tzer_cov_bits = clipper(tzer_cov) 39 | 40 | nnsmith_cov_bits = None 41 | for file in os.listdir(args.nnsmith): 42 | if file.endswith(".memcov.pkl"): 43 | with open(os.path.join(args.nnsmith, file), "rb") as fp: 44 | nnsmith_cov = pickle.load(fp) 45 | tmp_nnsmith_cov_bits = clipper(nnsmith_cov) 46 | if nnsmith_cov_bits is None: 47 | nnsmith_cov_bits = tmp_nnsmith_cov_bits 48 | else: 49 | nnsmith_cov_bits = np.logical_or(nnsmith_cov_bits, tmp_nnsmith_cov_bits) 50 | 51 | print(f"Tzer Memcov: {np.count_nonzero(tzer_cov_bits)}") 52 | print(f"NNSmith Memcov: {np.count_nonzero(nnsmith_cov_bits)}") 53 | tzer_unique = np.count_nonzero( 54 | np.logical_and( 55 | tzer_cov_bits, 56 | np.logical_not(np.logical_and(tzer_cov_bits, nnsmith_cov_bits)), 57 | ) 58 | ) 59 | nnsmith_unique = np.count_nonzero( 60 | np.logical_and( 61 | nnsmith_cov_bits, 62 | np.logical_not(np.logical_and(tzer_cov_bits, nnsmith_cov_bits)), 63 | ) 64 | ) 65 | print(f"Tzer Unique: {tzer_unique}") 66 | print(f"NNSmith Unique: {nnsmith_unique}") 67 | print( 68 | f"Common: {np.count_nonzero(np.logical_and(tzer_cov_bits, nnsmith_cov_bits))}" 69 | ) 70 | -------------------------------------------------------------------------------- /tests/torch/test_dump_load.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | import torch 4 | 5 | from nnsmith.graph_gen import model_gen 6 | from nnsmith.materialize import Model, Oracle, TestCase 7 | from nnsmith.narrow_spec import auto_opset 8 | 9 | TestCase.__test__ = False # supress PyTest warning 10 | 11 | 12 | def compare_two_oracle(src: Oracle, loaded: Oracle): 13 | assert len(loaded.input) == len(src.input) 14 | assert len(loaded.output) == len(src.output) 15 | for k, v in loaded.input.items(): 16 | assert np.allclose(v, src.input[k], equal_nan=True) 17 | for k, v in loaded.output.items(): 18 | assert np.allclose(v, src.output[k], equal_nan=True) 19 | 20 | 21 | def test_onnx_load_dump(tmp_path): 22 | d = tmp_path / "test_onnx_load_dump" 23 | d.mkdir() 24 | 25 | ONNXModelCPU = Model.init("onnx") 26 | 27 | gen = model_gen( 28 | opset=auto_opset(ONNXModelCPU), 29 | seed=54341, 30 | max_nodes=5, 31 | ) 32 | 33 | model = ONNXModelCPU.from_gir(gen.make_concrete()) 34 | 35 | assert model.with_torch 36 | 37 | model.refine_weights() # either random generated or gradient-based. 38 | oracle = model.make_oracle() 39 | 40 | testcase = TestCase(model, oracle) 41 | testcase.dump(root_folder=d) 42 | 43 | loaded_testcase = TestCase.load(model_type=type(model), root_folder=d) 44 | 45 | # check oracle 46 | compare_two_oracle(oracle, loaded_testcase.oracle) 47 | 48 | loaded_model = loaded_testcase.model.torch_model 49 | loaded_model.sat_inputs = {k: torch.from_numpy(v) for k, v in oracle.input.items()} 50 | rerun_oracle = loaded_model.make_oracle() 51 | compare_two_oracle(oracle, rerun_oracle) 52 | 53 | 54 | def test_bug_report_load_dump(tmp_path): 55 | d = tmp_path / "test_onnx_load_dump" 56 | d.mkdir() 57 | 58 | ONNXModelCPU = Model.init("onnx") 59 | gen = model_gen( 60 | opset=auto_opset(ONNXModelCPU), 61 | seed=5341, 62 | max_nodes=5, 63 | ) 64 | 65 | model = ONNXModelCPU.from_gir(gen.make_concrete()) 66 | 67 | assert model.with_torch 68 | 69 | model.refine_weights() # either random generated or gradient-based. 70 | oracle = model.make_oracle() 71 | 72 | testcase = TestCase(model, oracle) 73 | testcase.dump(root_folder=d) 74 | 75 | loaded_testcase = TestCase.load(model_type=type(model), root_folder=d) 76 | 77 | # check oracle 78 | compare_two_oracle(oracle, loaded_testcase.oracle) 79 | 80 | loaded_model = loaded_testcase.model.torch_model 81 | loaded_model.sat_inputs = {k: torch.from_numpy(v) for k, v in oracle.input.items()} 82 | rerun_oracle = loaded_model.make_oracle() 83 | compare_two_oracle(oracle, rerun_oracle) 84 | -------------------------------------------------------------------------------- /nnsmith/filter.py: -------------------------------------------------------------------------------- 1 | from inspect import signature 2 | from types import FunctionType 3 | from typing import Set, Type 4 | 5 | from nnsmith.materialize import BugReport, Stage, Symptom 6 | 7 | FILTERS = {} 8 | 9 | 10 | class filter: 11 | def __init__(self, name): 12 | self.name = name 13 | 14 | def __call__(self, fn_or_cls): 15 | assert self.name not in FILTERS, f"Filter {self.name} already exists." 16 | if isinstance(fn_or_cls, Type): # Class checks 17 | assert not signature( 18 | fn_or_cls 19 | ).parameters, f"filter class {fn_or_cls.__name__} (aka {self.name}) should not have any parameters." 20 | caller_sig = signature(fn_or_cls.__call__).parameters 21 | assert ( 22 | caller_sig["self"] and len(caller_sig) == 2 23 | ), f"filter class {fn_or_cls.__name__} (aka {self.name}) should implement __call__(self, report: BugReport)." 24 | elif isinstance(fn_or_cls, FunctionType): # Function checks 25 | caller_sig = signature(fn_or_cls).parameters 26 | assert ( 27 | len(caller_sig) == 1 28 | ), f"filter function {fn_or_cls.__name__} (aka {self.name}) should implement __call__(report: BugReport)." 29 | else: 30 | raise ValueError( 31 | f"filter {fn_or_cls} (aka {self.name}) should be a class or function." 32 | ) 33 | 34 | FILTERS[self.name] = fn_or_cls 35 | return fn_or_cls 36 | 37 | 38 | @filter("nan") 39 | def filter_nan(report: BugReport) -> bool: # True means filter; 40 | if report.symptom != Symptom.INCONSISTENCY or report.stage != Stage.VERIFICATION: 41 | return False 42 | 43 | # numpy.assert_allclose style. 44 | # TODO(ganler): can we use more well-formed checking? say directly checking the results? 45 | return ( 46 | "nan location mismatch" in report.log 47 | or "-9223372036854775808" in report.log # tf.cast(nan, int) is UB. 48 | or "-2147483648" in report.log 49 | ) 50 | 51 | 52 | @filter("inf") 53 | def filter_inf(report: BugReport) -> bool: 54 | if report.symptom != Symptom.INCONSISTENCY or report.stage != Stage.VERIFICATION: 55 | return False 56 | 57 | # numpy.assert_allclose style. 58 | return "inf" in report.log.replace("Max relative difference: inf", "") 59 | 60 | 61 | @filter("dup") # duplicate 62 | class FilterDup: 63 | def __init__(self): 64 | self.seen: Set[int] = set() 65 | 66 | def __call__(self, report: BugReport) -> bool: 67 | if ( 68 | report.symptom != Symptom.EXCEPTION 69 | and report.symptom != Symptom.INCONSISTENCY 70 | ): 71 | return False # don't filter bugs other than inconsistency/exception 72 | 73 | str_hash = hash(report.log) 74 | if str_hash in self.seen: 75 | return True 76 | 77 | self.seen.add(str_hash) 78 | return False 79 | 80 | 81 | # You can patch your own filters! 82 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .history 2 | 3 | # NNSmith-specific 4 | *.onnx 5 | tmp/ 6 | *.png 7 | *.pdf 8 | *.dot 9 | *.pkl 10 | *.stderr 11 | *.stdout 12 | *.lcov 13 | *.txt 14 | *.lz4 15 | *.profraw 16 | fuzz_report/ 17 | nnsmith_output/ 18 | nnsmith/_version.py 19 | 20 | # hydra 21 | outputs/ 22 | 23 | # Byte-compiled / optimized / DLL files 24 | __pycache__/ 25 | *.py[cod] 26 | *$py.class 27 | *.pkl_memoize_py3 28 | 29 | # C extensions 30 | # *.so 31 | 32 | # Distribution / packaging 33 | .Python 34 | build/ 35 | develop-eggs/ 36 | dist/ 37 | downloads/ 38 | eggs/ 39 | .eggs/ 40 | lib/ 41 | lib64/ 42 | parts/ 43 | sdist/ 44 | var/ 45 | wheels/ 46 | share/python-wheels/ 47 | *.egg-info/ 48 | .installed.cfg 49 | *.egg 50 | MANIFEST 51 | 52 | # PyInstaller 53 | # Usually these files are written by a python script from a template 54 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 55 | *.manifest 56 | *.spec 57 | 58 | # Installer logs 59 | pip-log.txt 60 | pip-delete-this-directory.txt 61 | 62 | # Unit test / coverage reports 63 | htmlcov/ 64 | .tox/ 65 | .nox/ 66 | .coverage 67 | .coverage.* 68 | .cache 69 | nosetests.xml 70 | coverage.xml 71 | *.cover 72 | *.py,cover 73 | .hypothesis/ 74 | .pytest_cache/ 75 | cover/ 76 | 77 | # Translations 78 | *.mo 79 | *.pot 80 | 81 | # Django stuff: 82 | *.log 83 | local_settings.py 84 | db.sqlite3 85 | db.sqlite3-journal 86 | 87 | # Flask stuff: 88 | instance/ 89 | .webassets-cache 90 | 91 | # Scrapy stuff: 92 | .scrapy 93 | 94 | # Sphinx documentation 95 | docs/_build/ 96 | 97 | # PyBuilder 98 | .pybuilder/ 99 | target/ 100 | 101 | # Jupyter Notebook 102 | .ipynb_checkpoints 103 | 104 | # IPython 105 | profile_default/ 106 | ipython_config.py 107 | 108 | # pyenv 109 | # For a library or package, you might want to ignore these files since the code is 110 | # intended to run in multiple environments; otherwise, check them in: 111 | # .python-version 112 | 113 | # pipenv 114 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 115 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 116 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 117 | # install all needed dependencies. 118 | #Pipfile.lock 119 | 120 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 121 | __pypackages__/ 122 | 123 | # Celery stuff 124 | celerybeat-schedule 125 | celerybeat.pid 126 | 127 | # SageMath parsed files 128 | *.sage.py 129 | 130 | # Environments 131 | .vscode 132 | .env 133 | .venv 134 | env/ 135 | venv/ 136 | ENV/ 137 | env.bak/ 138 | venv.bak/ 139 | 140 | # Spyder project settings 141 | .spyderproject 142 | .spyproject 143 | 144 | # Rope project settings 145 | .ropeproject 146 | 147 | # mkdocs documentation 148 | /site 149 | 150 | # mypy 151 | .mypy_cache/ 152 | .dmypy.json 153 | dmypy.json 154 | 155 | # Pyre type checker 156 | .pyre/ 157 | 158 | # pytype static type analyzer 159 | .pytype/ 160 | 161 | # Cython debug symbols 162 | cython_debug/ 163 | 164 | # Exclude 165 | !requirements/**/*.txt 166 | .DS_Store 167 | -------------------------------------------------------------------------------- /experiments/evaluate_models.py: -------------------------------------------------------------------------------- 1 | """ 2 | Given the directory containing all tests. We replay the test execution and record coverage in LLVM profraw format. 3 | The intermediate tests can be saved using fuzz.save_test={{DIR_TO_SAVE}}. 4 | """ 5 | import multiprocessing as mp 6 | import os 7 | import subprocess 8 | 9 | from nnsmith.util import mkdir 10 | 11 | 12 | def batched(iterable, n=1): 13 | l = len(iterable) 14 | for ndx in range(0, l, n): 15 | yield iterable[ndx : min(ndx + n, l)] 16 | 17 | 18 | def model_exec(test_paths, model_type, backend_type, backend_target, profraw_path): 19 | model_paths = [] 20 | for test_path in test_paths: 21 | for file in os.listdir(test_path): 22 | if file.startswith("model"): 23 | model_paths.append(os.path.join(test_path, file)) 24 | break 25 | 26 | arguments = [ 27 | "python3", 28 | "nnsmith/cli/model_exec.py", 29 | "model.type=" + model_type, 30 | "backend.type=" + backend_type, 31 | "backend.target=" + backend_target, 32 | f"model.path={model_paths}", 33 | ] 34 | 35 | copied_env = os.environ.copy() 36 | copied_env["LLVM_PROFILE_FILE"] = profraw_path 37 | 38 | p = subprocess.Popen( 39 | arguments, # Show all output 40 | env=copied_env, 41 | ) 42 | p.communicate() 43 | exit_code = p.returncode 44 | 45 | if exit_code != 0: 46 | print( 47 | f"==> model_exec crashed when generating {profraw_path}! => EXIT CODE {exit_code}" 48 | ) 49 | 50 | 51 | if __name__ == "__main__": 52 | import argparse 53 | 54 | parser = argparse.ArgumentParser() 55 | parser.add_argument( 56 | "--root", type=str, required=True, help="Folder to all the tests." 57 | ) 58 | parser.add_argument("--batch-size", type=int, default=100, help="") 59 | parser.add_argument("--model_type", type=str, required=True, help="Model type.") 60 | parser.add_argument("--backend_type", type=str, required=True, help="Backend type.") 61 | parser.add_argument( 62 | "--backend_target", type=str, required=True, help="Say `cpu` or `cuda`." 63 | ) 64 | parser.add_argument( 65 | "--parallel", type=int, default=8, help="Number of process for execution." 66 | ) 67 | 68 | args = parser.parse_args() 69 | 70 | time2path = {} 71 | for dir in os.listdir(args.root): 72 | if dir != "coverage": 73 | time2path[float(dir)] = os.path.join(args.root, dir) 74 | 75 | time_stamps = sorted(time2path.keys()) 76 | batches = list(batched(time_stamps, args.batch_size)) 77 | 78 | print(f"=> Number of batches: {len(batches)} of size {args.batch_size}") 79 | 80 | cov_save = os.path.join(args.root, "coverage") 81 | 82 | mkdir(cov_save) 83 | 84 | def batch_exec(batch): 85 | batch_paths = [time2path[time] for time in batch] 86 | profraw_path = os.path.join(cov_save, f"{max(batch)}.profraw") 87 | model_exec( 88 | batch_paths, 89 | args.model_type, 90 | args.backend_type, 91 | args.backend_target, 92 | profraw_path, 93 | ) 94 | 95 | with mp.Pool(processes=args.parallel) as pool: 96 | pool.map(batch_exec, batches) 97 | -------------------------------------------------------------------------------- /exp_helper/4_throughput_raw.py: -------------------------------------------------------------------------------- 1 | import re 2 | import pandas as pd 3 | import sys 4 | import os 5 | 6 | TIME_LOG_DIR_PATH = "/root/raw_data/raw_data/tab5_raw_data" 7 | OUTPUT_DIR_PATH = "/root/raw_data/raw_data_results" 8 | 9 | 10 | def print_avg_time(file_path): 11 | pattern = re.compile( 12 | r"Timing: gen: \s*(-?\d+\.\d+?)ms, saturation: \s*(-?\d+\.\d+?)ms, egraph_gen: \s*(-?\d+\.\d+?)ms, eval: \s*(-?\d+\.\d+?)ms,") 13 | record = [] 14 | with open(file_path, "r") as f: 15 | for line in f.readlines(): 16 | match = pattern.search(line) 17 | if match is not None: 18 | assert len(match.groups()) == 4 19 | tmp = [] 20 | for i in range(4): 21 | tmp.append(float(match.group(i + 1))) 22 | record.append(tmp) 23 | df = pd.DataFrame(data=record, columns=["gen", "rewriting", "extraction", "execution"]) 24 | return df 25 | 26 | 27 | def main(): 28 | df_xla_5 = print_avg_time(os.path.join(TIME_LOG_DIR_PATH, "./time_log_xla_5")) 29 | df_xla_10 = print_avg_time(os.path.join(TIME_LOG_DIR_PATH, "./time_log_xla_10")) 30 | df_xla_15 = print_avg_time(os.path.join(TIME_LOG_DIR_PATH, "./time_log_xla_15")) 31 | df_tvm_5 = print_avg_time(os.path.join(TIME_LOG_DIR_PATH, "./time_log_tvm_5")) 32 | df_tvm_10 = print_avg_time(os.path.join(TIME_LOG_DIR_PATH, "./time_log_tvm_10")) 33 | df_tvm_15 = print_avg_time(os.path.join(TIME_LOG_DIR_PATH, "./time_log_tvm_15")) 34 | 35 | df = pd.DataFrame(columns=["compiler", "node_num", "gen", "rewriting", "extraction", "execution"]) 36 | 37 | df = df.append( 38 | {"compiler": "tvm", "node_num": 5, "gen": df_tvm_5["gen"].mean(), "rewriting": df_tvm_5["rewriting"].mean(), 39 | "extraction": df_tvm_5["extraction"].mean(), "execution": df_tvm_5["execution"].mean()}, ignore_index=True) 40 | df = df.append( 41 | {"compiler": "tvm", "node_num": 10, "gen": df_tvm_10["gen"].mean(), "rewriting": df_tvm_10["rewriting"].mean(), 42 | "extraction": df_tvm_10["extraction"].mean(), "execution": df_tvm_10["execution"].mean()}, ignore_index=True) 43 | df = df.append( 44 | {"compiler": "tvm", "node_num": 15, "gen": df_tvm_15["gen"].mean(), "rewriting": df_tvm_15["rewriting"].mean(), 45 | "extraction": df_tvm_15["extraction"].mean(), "execution": df_tvm_15["execution"].mean()}, ignore_index=True) 46 | df = df.append( 47 | {"compiler": "xla", "node_num": 5, "gen": df_xla_5["gen"].mean(), "rewriting": df_xla_5["rewriting"].mean(), 48 | "extraction": df_xla_5["extraction"].mean(), "execution": df_xla_5["execution"].mean()}, ignore_index=True) 49 | df = df.append( 50 | {"compiler": "xla", "node_num": 10, "gen": df_xla_10["gen"].mean(), "rewriting": df_xla_10["rewriting"].mean(), 51 | "extraction": df_xla_10["extraction"].mean(), "execution": df_xla_10["execution"].mean()}, ignore_index=True) 52 | df = df.append( 53 | {"compiler": "xla", "node_num": 15, "gen": df_xla_15["gen"].mean(), "rewriting": df_xla_15["rewriting"].mean(), 54 | "extraction": df_xla_15["extraction"].mean(), "execution": df_xla_15["execution"].mean()}, ignore_index=True) 55 | 56 | print(df) 57 | 58 | output_path = os.path.join(OUTPUT_DIR_PATH, "tab5.csv") 59 | if os.path.exists(output_path): 60 | os.remove(output_path) 61 | df.to_csv(output_path, index=False) 62 | 63 | 64 | if __name__ == '__main__': 65 | main() 66 | -------------------------------------------------------------------------------- /nnsmith/materialize/torch/dialect.py: -------------------------------------------------------------------------------- 1 | from math import prod 2 | from random import randint 3 | from typing import List, Tuple, Union 4 | 5 | from nnsmith.abstract.arith import * 6 | from nnsmith.abstract.dtype import ( 7 | DTYPE_GEN_ALL, 8 | DTYPE_GEN_COMPLEX, 9 | DTYPE_GEN_FLOATS, 10 | DTYPE_GEN_INTS, 11 | DTYPE_GEN_NON_BOOL, 12 | DType, 13 | ) 14 | from nnsmith.abstract.op import ( 15 | MatMul, 16 | ReduceBase, 17 | UnaryOpBase, 18 | mark_materialize, 19 | rank_from, 20 | ) 21 | from nnsmith.abstract.tensor import AbsTensor 22 | from nnsmith.error import ConstraintCheck 23 | 24 | 25 | @mark_materialize("torch") 26 | class Linear(UnaryOpBase): 27 | in_dtypes = [(DType.float32,)] 28 | out_dtypes = [(DType.float32,)] 29 | 30 | def __init__(self, ifeat: Union[int, z3.ExprRef], ofeat: Union[int, z3.ExprRef]): 31 | super().__init__() 32 | self.ifeat = ifeat 33 | self.ofeat = ofeat 34 | self.inp_ranks = [rank_from(1)] 35 | self.out_ranks = [rank_from(1)] 36 | 37 | def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]: 38 | assert len(input_shapes) == 1, "Linear only takes one input, but got {}".format( 39 | len(input_shapes) 40 | ) 41 | return [ 42 | AbsTensor( 43 | shape=[*input_shapes[0].shape[:-1], self.ofeat], dtype=DType.float32 44 | ) 45 | ] 46 | 47 | def requires(self, input_shapes: List[AbsTensor]) -> List[z3.ExprRef]: 48 | ConstraintCheck.true(input_shapes[0].ndims >= 1) 49 | return [ 50 | nnsmith_ge(self.ifeat, 1), 51 | nnsmith_ge(self.ofeat, 1), 52 | nnsmith_eq(input_shapes[0].shape[-1], self.ifeat), 53 | ] 54 | 55 | def deduct_inp_ranks_and_dtype( 56 | self, out_abs_tensor: List[AbsTensor] 57 | ) -> List[Tuple[int, DType]]: 58 | return [(out_abs_tensor[0].ndims, DType.float32)] 59 | 60 | 61 | @mark_materialize("torch") 62 | class Flatten(UnaryOpBase): 63 | in_dtypes = [(i,) for i in DTYPE_GEN_ALL] 64 | out_dtypes = [(i,) for i in DTYPE_GEN_ALL] 65 | 66 | def __init__(self): 67 | super().__init__() 68 | self.inp_ranks = [rank_from(1)] 69 | self.out_ranks = [(1,)] 70 | 71 | def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]: 72 | inp = input_shapes[0] 73 | return [ 74 | AbsTensor( 75 | shape=[prod(inp.shape)], 76 | dtype=inp.dtype, 77 | ) 78 | ] 79 | 80 | def requires(self, input_shapes): 81 | return [] 82 | 83 | def deduct_inp_ranks_and_dtype( 84 | self, out_abs_tensor: List[AbsTensor] 85 | ) -> List[Tuple[int, DType]]: 86 | return [(randint(0, 4), out_abs_tensor[0].dtype)] 87 | 88 | 89 | @mark_materialize("torch") 90 | class TorchReduceSum(ReduceBase): 91 | in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL] 92 | out_dtypes = [(i,) for i in DTYPE_GEN_FLOATS + DTYPE_GEN_COMPLEX + [DType.int64]] 93 | 94 | def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]: 95 | output = super().type_transfer(input_shapes) 96 | if input_shapes[0].dtype in DTYPE_GEN_INTS: # This is a PyTorch trick... 97 | output[0].dtype = DType.int64 98 | return output 99 | 100 | 101 | @mark_materialize("torch") 102 | class PTMatMul(MatMul): 103 | pass 104 | -------------------------------------------------------------------------------- /nnsmith/abstract/tensor.py: -------------------------------------------------------------------------------- 1 | from functools import reduce 2 | from typing import List, Union 3 | 4 | import z3 5 | 6 | from nnsmith.abstract.arith import * 7 | from nnsmith.abstract.dtype import DType 8 | from nnsmith.error import ConstraintCheck, SanityCheck 9 | 10 | 11 | class AbsTensor: 12 | def __init__(self, shape: List[Union[int, z3.ExprRef]], dtype: DType): 13 | assert isinstance( 14 | shape, (list, tuple) 15 | ), f"Shape must be a list/tuple, but got {shape}" 16 | self.shape = list(shape) 17 | self.dtype = DType(dtype) 18 | 19 | def downcast_rank(self): 20 | return AbsTensor(shape=[None] * self.ndims, dtype=self.dtype) 21 | 22 | def __hash__(self) -> int: 23 | return hash((tuple(self.shape), self.dtype)) 24 | 25 | def __repr__(self) -> str: 26 | return f"AbsTensor<{self.dtype.short()}>{str(self.shape)}" 27 | 28 | def pretty(self) -> str: 29 | return f"{self.dtype.short()}{self.shape}" 30 | 31 | def weak_compare(self, other: "AbsTensor") -> bool: 32 | if self.dtype != other.dtype or self.ndims != other.ndims: 33 | return False 34 | for l, r in zip(self.shape, other.shape): 35 | if isinstance(l, z3.ExprRef) or isinstance(r, z3.ExprRef): 36 | continue 37 | if l != r: 38 | return False 39 | return True 40 | 41 | def strong_compare(self, other: "AbsTensor") -> bool: 42 | return self.shape == other.shape and self.dtype == other.dtype 43 | 44 | def __eq__(self, other: "AbsTensor") -> bool: 45 | return self.strong_compare(other) 46 | 47 | def ge_zero(self): 48 | ret = [] 49 | for s in self.shape: 50 | if isinstance(s, z3.ExprRef): 51 | ret.append(nnsmith_ge(s, 0)) 52 | else: 53 | ConstraintCheck.ge(s, 0) 54 | return ret 55 | 56 | def sym_gt_conc_ge_zero(self): 57 | ret = [] 58 | for s in self.shape: 59 | if isinstance(s, z3.ExprRef): 60 | ret.append(nnsmith_gt(s, 0)) 61 | else: 62 | ConstraintCheck.ge(s, 0) 63 | return ret 64 | 65 | def gt_zero(self): 66 | ret = [] 67 | for s in self.shape: 68 | if isinstance(s, z3.ExprRef): 69 | ret.append(nnsmith_gt(s, 0)) 70 | else: 71 | ConstraintCheck.gt(s, 0) 72 | return ret 73 | 74 | def eq(self, other): 75 | SanityCheck.eq(self.ndims, other.ndims) 76 | ret = [] 77 | for i in range(self.ndims): 78 | if isinstance(self.shape[i], z3.ExprRef) or isinstance( 79 | other.shape[i], z3.ExprRef 80 | ): 81 | ret.append(nnsmith_eq(self.shape[i], other.shape[i])) 82 | else: 83 | ConstraintCheck.eq(self.shape[i], other.shape[i]) 84 | return ret 85 | 86 | def torch(self): 87 | import torch 88 | 89 | return torch.Size(self.shape) 90 | 91 | def constains_symbol(self) -> bool: 92 | return any(isinstance(s, z3.ExprRef) for s in self.shape) 93 | 94 | def nelement(self): 95 | if len(self.shape) == 0: # Scalar 96 | return 1 97 | return reduce(lambda x, y: nnsmith_mul(x, y), self.shape, 1) 98 | 99 | def nbytes(self) -> int: 100 | return self.nelement() * self.dtype.sizeof() 101 | 102 | def deepcopy(self): 103 | return AbsTensor(shape=list(self.shape), dtype=self.dtype) 104 | 105 | @property 106 | def ndims(self): 107 | return len(self.shape) 108 | 109 | def is_concrete(self) -> bool: 110 | return all(isinstance(s, int) for s in self.shape) 111 | 112 | def htype(self): # High-level type 113 | return (self.dtype, self.ndims) 114 | -------------------------------------------------------------------------------- /nnsmith/materialize/torch/proxy_grad.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from multipledispatch import dispatch 3 | 4 | from nnsmith.abstract.op import * 5 | 6 | # used proxy gradient functions 7 | SLOPE = 0.1 8 | 9 | 10 | class PGTruncFuncBase(torch.autograd.Function): 11 | # incomplete class 12 | @staticmethod 13 | def backward(ctx, grad_output): 14 | # let f' = x 15 | return grad_output 16 | 17 | 18 | class PGFloorFunc(PGTruncFuncBase): 19 | @staticmethod 20 | def forward(ctx, input): 21 | return torch.floor(input) 22 | 23 | 24 | class PGRoundFunc(PGTruncFuncBase): 25 | @staticmethod 26 | def forward(ctx, input): 27 | return torch.round(input) 28 | 29 | 30 | class PGCeilFunc(PGTruncFuncBase): 31 | @staticmethod 32 | def forward(ctx, input): 33 | return torch.ceil(input) 34 | 35 | 36 | class PGReLUFunc(torch.autograd.Function): 37 | @staticmethod 38 | def forward(ctx, input): 39 | ctx.save_for_backward(input) 40 | if input.dtype == torch.float16: 41 | return torch.relu(input.float()).half() 42 | return torch.relu(input) 43 | 44 | @staticmethod 45 | def backward(ctx, grad_output): 46 | (input,) = ctx.saved_tensors 47 | # let f' = l * x in bound 48 | # let f' = x out bound 49 | grad_input = grad_output.clone() 50 | return torch.where(input > 0, grad_input, grad_input * SLOPE) 51 | 52 | 53 | class PGClipFunc(torch.autograd.Function): 54 | @staticmethod 55 | def forward(ctx, input, min, max): 56 | ctx.save_for_backward(input) 57 | ctx.clip_min = min 58 | ctx.clip_max = max 59 | return torch.clip(input, min=min, max=max) 60 | 61 | @staticmethod 62 | def backward(ctx, grad_output): 63 | (input,) = ctx.saved_tensors 64 | # let f' = l * x in bound 65 | # let f' = x out bound 66 | grad_input = grad_output.clone() 67 | return ( 68 | torch.where( 69 | (input > ctx.clip_max).logical_or(input < ctx.clip_min), 70 | grad_input * SLOPE, 71 | grad_input, 72 | ), 73 | None, 74 | None, 75 | ) 76 | 77 | 78 | # Modules. 79 | 80 | 81 | class PGCeil(torch.nn.Module): 82 | def __init__(self): 83 | super(PGCeil, self).__init__() 84 | 85 | def forward(self, x): 86 | return PGCeilFunc.apply(x) 87 | 88 | 89 | class PGFloor(torch.nn.Module): 90 | def __init__(self): 91 | super(PGFloor, self).__init__() 92 | 93 | def forward(self, x): 94 | return PGFloorFunc.apply(x) 95 | 96 | 97 | class PGRound(torch.nn.Module): 98 | def __init__(self): 99 | super(PGRound, self).__init__() 100 | 101 | def forward(self, x): 102 | return PGRoundFunc.apply(x) 103 | 104 | 105 | class PGReLU(torch.nn.Module): 106 | def __init__(self): 107 | super(PGReLU, self).__init__() 108 | 109 | def forward(self, x): 110 | return PGReLUFunc.apply(x) 111 | 112 | 113 | class PGClip(torch.nn.Module): 114 | def __init__(self, min, max): 115 | self.min = min 116 | self.max = max 117 | super(PGClip, self).__init__() 118 | 119 | def forward(self, x): 120 | return PGClipFunc.apply(x, self.min, self.max) 121 | 122 | 123 | # proxy_fn: proxy 124 | 125 | 126 | @dispatch(ReLU) 127 | def proxy_fn(op: ReLU): 128 | return PGReLU() 129 | 130 | 131 | @dispatch(Ceil) 132 | def proxy_fn(op: Ceil): 133 | return PGCeil() 134 | 135 | 136 | # PGFloor 137 | @dispatch(Floor) 138 | def proxy_fn(op: Floor): 139 | return PGFloor() 140 | 141 | 142 | # PGClip 143 | @dispatch(Clip) 144 | def proxy_fn(op: Clip): 145 | if op.input_like[0].dtype in DTYPE_GEN_FLOATS: 146 | return PGClip(-1.5, 1.5) 147 | else: 148 | return PGClip(-1, 1) 149 | 150 | 151 | # Round 152 | @dispatch(Round) 153 | def proxy_fn(op: Round): 154 | return PGRound() 155 | -------------------------------------------------------------------------------- /exp_helper/4_run_throughput.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import shutil 3 | import os 4 | import copy 5 | import re 6 | import pandas as pd 7 | 8 | POLYJUICE_ROOT_PATH = "/root/polyjuice/" 9 | OUTPUT_DIR = "/root/4_throughput/" 10 | 11 | TMP_DIR = "/tmp/throughput/" 12 | 13 | MAX_TIME = 2 * 60 * 60 14 | 15 | 16 | def print_avg_time(file_path): 17 | pattern = re.compile( 18 | r"Timing: gen: \s*(-?\d+\.\d+?)ms, saturation: \s*(-?\d+\.\d+?)ms, egraph_gen: \s*(-?\d+\.\d+?)ms, eval: \s*(-?\d+\.\d+?)ms,") 19 | record = [] 20 | with open(file_path, "r") as f: 21 | for line in f.readlines(): 22 | match = pattern.search(line) 23 | if match is not None: 24 | assert len(match.groups()) == 4 25 | tmp = [] 26 | for i in range(4): 27 | tmp.append(float(match.group(i + 1))) 28 | record.append(tmp) 29 | df = pd.DataFrame(data=record, columns=["gen", "rewriting", "extraction", "execution"]) 30 | return df 31 | 32 | 33 | def main(): 34 | if os.path.exists(OUTPUT_DIR): 35 | shutil.rmtree(OUTPUT_DIR) 36 | os.makedirs(OUTPUT_DIR, exist_ok=True) 37 | 38 | if os.path.exists(TMP_DIR): 39 | shutil.rmtree(TMP_DIR) 40 | os.makedirs(TMP_DIR, exist_ok=True) 41 | 42 | new_env = copy.deepcopy(os.environ) 43 | new_env["PYTHONPATH"] = "/root/polyjuice" 44 | new_env["EQ_HELPER_PATH"] = "/root/polyjuice/libdl_compiler_fuzzer_helper.so" 45 | new_env["SAVE_EQ_IS_WRONG"] = "false" 46 | new_env["TIME_EXP_RUN"] = "true" 47 | 48 | tvm_5_log_path = f"{OUTPUT_DIR}/time_log_tvm_5" 49 | tvm_10_log_path = f"{OUTPUT_DIR}/time_log_tvm_10" 50 | 51 | print("start running") 52 | tvm_5_f = open(tvm_5_log_path, "w") 53 | tvm_10_f = open(tvm_10_log_path, "w") 54 | tvm_5_node_proc = subprocess.Popen( 55 | ["python3", f"{POLYJUICE_ROOT_PATH}/nnsmith/cli/equivalent_fuzz.py", f"fuzz.time={MAX_TIME}", 56 | "model.type=onnx", 57 | "backend.type=tvm", 58 | "backend.target=cpu", f"fuzz.root={TMP_DIR}/tvm_5", "debug.viz=true", "mgen.max_nodes=5"], env=new_env, 59 | cwd=POLYJUICE_ROOT_PATH, stdout=tvm_5_f, stderr=tvm_5_f) 60 | tvm_10_node_proc = subprocess.Popen( 61 | ["python3", f"{POLYJUICE_ROOT_PATH}/nnsmith/cli/equivalent_fuzz.py", f"fuzz.time={MAX_TIME}", 62 | "model.type=onnx", 63 | "backend.type=tvm", 64 | "backend.target=cpu", f"fuzz.root={TMP_DIR}/tvm_10", "debug.viz=true", "mgen.max_nodes=10"], 65 | env=new_env, 66 | cwd=POLYJUICE_ROOT_PATH, stdout=tvm_10_f, stderr=tvm_10_f) 67 | try: 68 | tvm_5_node_proc.wait(timeout=MAX_TIME + 60) 69 | except subprocess.TimeoutExpired: 70 | tvm_5_node_proc.kill() 71 | print("tvm_5 killed because of timeout") 72 | 73 | try: 74 | tvm_10_node_proc.wait(timeout=MAX_TIME + 60) 75 | except subprocess.TimeoutExpired: 76 | tvm_10_node_proc.kill() 77 | print("tvm_10 killed because of timeout") 78 | 79 | tvm_5_f.close() 80 | tvm_10_f.close() 81 | 82 | print("finished running") 83 | 84 | df_tvm_5 = print_avg_time(tvm_5_log_path) 85 | df_tvm_10 = print_avg_time(tvm_10_log_path) 86 | 87 | df = pd.DataFrame(columns=["compiler", "node_num", "gen", "rewriting", "extraction", "execution"]) 88 | 89 | df = df.append( 90 | {"compiler": "tvm", "node_num": 5, "gen": df_tvm_5["gen"].mean(), "rewriting": df_tvm_5["rewriting"].mean(), 91 | "extraction": df_tvm_5["extraction"].mean(), "execution": df_tvm_5["execution"].mean()}, ignore_index=True) 92 | df = df.append( 93 | {"compiler": "tvm", "node_num": 10, "gen": df_tvm_10["gen"].mean(), "rewriting": df_tvm_10["rewriting"].mean(), 94 | "extraction": df_tvm_10["extraction"].mean(), "execution": df_tvm_10["execution"].mean()}, ignore_index=True) 95 | 96 | print(df) 97 | df.to_csv(f"{OUTPUT_DIR}/tab5.csv", index=False) 98 | 99 | 100 | if __name__ == '__main__': 101 | main() 102 | -------------------------------------------------------------------------------- /nnsmith/materialize/tensorflow/tfnet.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import Callable, Dict, List, cast 3 | 4 | import tensorflow as tf 5 | 6 | from nnsmith.abstract.op import AbsOpBase, Input, Constant 7 | from nnsmith.error import SanityCheck 8 | from nnsmith.gir import GraphIR 9 | from nnsmith.logging import TF_LOG 10 | from nnsmith.materialize.tensorflow.forward import forward_fn 11 | 12 | 13 | @dataclass 14 | class Instr: 15 | fwd_fn: Callable 16 | inp_keys: List[str] 17 | out_keys: List[str] 18 | 19 | 20 | class TFNet(tf.Module): 21 | """A TensorFlow network whose computation is defined by a GraphIR.""" 22 | 23 | def __init__(self, ir: GraphIR) -> None: 24 | """Build a TensorFlow model from GraphIR 25 | Args: 26 | ir (GraphIR): minimal information for constructing a concrete graph. 27 | """ 28 | super().__init__() 29 | self.ir: GraphIR = ir 30 | self.mlist: List[Callable] = [] 31 | self.instructions: List[Instr] = [] 32 | self.special_op: Dict[int, tf.Module] = {} 33 | self.input_name_map: Dict[int, str] = {} 34 | self.constant_record: List[str] = [] 35 | 36 | for inst in self.ir.insts: 37 | if not isinstance(inst.iexpr.op, Input): 38 | op = cast(AbsOpBase, inst.iexpr.op) 39 | fwd_fn = forward_fn(op) 40 | SanityCheck.true(fwd_fn is not None, f"Bad impl for {inst.iexpr.op}") 41 | if isinstance(fwd_fn, tf.Module): 42 | self.mlist.append(fwd_fn) # Add tf.Module to track its parameters 43 | instruction_index = len(self.instructions) 44 | self.instructions.append(Instr(fwd_fn, inst.iexpr.args, inst.retvals())) 45 | 46 | if "op_index" in inst.iexpr.op.extra_attrs: 47 | self.special_op[inst.iexpr.op.extra_attrs["op_index"]] = instruction_index 48 | if isinstance(op, Constant): 49 | self.constant_record.append(inst.iexpr.op.extra_attrs["op_index"]) 50 | 51 | else: 52 | if "op_index" in inst.iexpr.op.extra_attrs: 53 | op_index = inst.iexpr.op.extra_attrs["op_index"] 54 | assert len(inst.retvals()) == 1 55 | self.input_name_map[op_index] = inst.retvals()[0] 56 | 57 | @tf.function(autograph=False) # disabling autograph makes it faster 58 | def __call__(self, *args, **kwargs) -> Dict[str, tf.Tensor]: 59 | return self.__forward(*args, **kwargs) 60 | 61 | @tf.function(autograph=False) 62 | def call_by_dict(self, x: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]: 63 | return self.__forward(**x) 64 | 65 | def __forward(self, *args, **kwargs) -> Dict[str, tf.Tensor]: 66 | mode = "Running Eagerly" if tf.executing_eagerly() else "Tracing" 67 | TF_LOG.debug(f"{mode} with JIT config: {tf.config.optimizer.get_jit()}") 68 | 69 | key2tensor: Dict[str, tf.Tensor] = {} 70 | if len(args) == len(self.ir.input_var()): 71 | for i, key in enumerate(self.ir.input_var()): 72 | key2tensor[key] = args[i] 73 | elif len(kwargs) == len(self.ir.input_var()): 74 | for i, key in enumerate(self.ir.input_var()): 75 | key2tensor[key] = kwargs[key] 76 | else: 77 | raise ValueError("Use either args or kwargs only") 78 | 79 | for instr in self.instructions: 80 | # get inputs 81 | inp_tensors = [key2tensor[key] for key in instr.inp_keys] 82 | 83 | # forward 84 | out_tensors = instr.fwd_fn(*inp_tensors) 85 | 86 | if isinstance(out_tensors, tf.Tensor): 87 | out_tensors = [out_tensors] 88 | if isinstance(out_tensors, tuple): 89 | out_tensors = list(out_tensors) 90 | 91 | # store outputs 92 | for i_out, out_key in enumerate(instr.out_keys): 93 | key2tensor[out_key] = out_tensors[i_out] 94 | 95 | return {k: key2tensor[k] for k in self.ir.leaf_var()} 96 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | ## About PolyJuice 2 | 3 | PolyJuice is a fuzzer that can generate equivalent computation graphs to test the correctness of tensor compilers. 4 | Its basic idea is to generate an equivalent tensor programs for a given tensor program, and then compare the output of 5 | the two programs to check the correctness of the compiler under test. 6 | 7 | ## Usage 8 | 9 | Before running the tool, we should first set some environment variables to enable the quality saturation engine: 10 | 11 | ```bash 12 | export PYTHONPATH="$PWD" 13 | export EQ_HELPER_PATH="$PWD/libdl_compiler_fuzzer_helper.so" 14 | ``` 15 | 16 | The code of this dynamic library is placed in the ``equality_saturation_helper/rust_helper``. 17 | 18 | Next, ``python3 nnsmith/cli/equivalent_fuzz.py --help`` shows how it works: 19 | 20 | ```text 21 | == Config == 22 | Override anything in the config (foo.bar=value) 23 | 24 | model: 25 | type: null 26 | path: ??? 27 | mgen: 28 | max_nodes: 5 29 | timeout_ms: 10000 30 | vulops: false 31 | method: symbolic-cinit 32 | save: nnsmith_output 33 | seed: null 34 | max_elem_per_tensor: 65536 35 | rank_choices: null 36 | dtype_choices: null 37 | include: null 38 | exclude: null 39 | patch_requires: [] 40 | grad_check: false 41 | backend: 42 | type: null 43 | optmax: true 44 | target: cpu 45 | ad: 46 | type: null 47 | cache: 48 | topset: true 49 | debug: 50 | viz: false 51 | viz_fmt: png 52 | fuzz: 53 | time: 14400 54 | root: ??? 55 | seed: null 56 | crash_safe: false 57 | test_timeout: null 58 | save_test: null 59 | random_num: 1 60 | filter: 61 | type: [] 62 | patch: [] 63 | cmp: 64 | equal_nan: true 65 | raw_input: null 66 | oracle: auto 67 | with: 68 | type: null 69 | optmax: true 70 | target: cpu 71 | seed: null 72 | bug_presence: report 73 | save: null 74 | 75 | 76 | Powered by Hydra (https://hydra.cc) 77 | Use --hydra-help to view Hydra specific help 78 | ``` 79 | 80 | Basically, the usage is almost identical with NNSmith. We provide some examples to show how to test Inductor, XLA and 81 | TVM with PolyJuice. 82 | 83 | ### Test Inductor 84 | 85 | To test PyTorch Inductor, make sure you have installed PyTorch (https://pytorch.org/get-started/locally/). We can run 86 | the following command to test Inductor: 87 | 88 | ```bash 89 | python3 nnsmith/cli/equivalent_fuzz.py fuzz.time=10s model.type=torch backend.type=pt2 backend.target=cpu fuzz.root=/tmp/fuzz_report debug.viz=true mgen.max_nodes=5 90 | ``` 91 | 92 | It will test Inductor for 10 seconds and output bugs to the directory ``/tmp/fuzz_report`` if any. 93 | 94 | ### Test XLA 95 | 96 | To test XLA, make sure you have installed TensorFlow (https://www.tensorflow.org/install). We can run the following 97 | command to test XLA: 98 | 99 | ```bash 100 | python3 nnsmith/cli/equivalent_fuzz.py fuzz.time=10s model.type=tensorflow backend.type=xla backend.target=cpu fuzz.root=/tmp/fuzz_report debug.viz=true mgen.max_nodes=5 101 | ``` 102 | 103 | It will test XLA for 10 seconds and output bugs to the directory ``/tmp/fuzz_report`` if any. 104 | 105 | ### Test TVM 106 | 107 | To test TVM, make sure you have installed TVM (https://tvm.apache.org/docs/install/index.html). We can run the following 108 | command to test TVM: 109 | 110 | ```bash 111 | python3 nnsmith/cli/equivalent_fuzz.py fuzz.time=10s model.type=onnx backend.type=tvm backend.target=cpu fuzz.root=/tmp/fuzz_report debug.viz=true mgen.max_nodes=5 112 | ``` 113 | 114 | It will test TVM for 10 seconds and output bugs to the directory ``/tmp/fuzz_report`` if any. 115 | 116 | ## Add New Backend 117 | 118 | To add a new backend, we need to implement the backend in the ``nnsmith/backend`` directory. There are some examples for 119 | adapting a new backend. You can refer to ``pt2.py`` and ``hidet.py``. 120 | 121 | ## Acknowledgement 122 | 123 | PolyJuice is built on the top of NNSmith, and relies on NNSmith's model generation to generate the initial tensor 124 | program. In addition, PolyJuice reuse equality saturation engine from egg. We would like to thank the authors of NNSmith 125 | and egg for their great work. -------------------------------------------------------------------------------- /nnsmith/backends/pt2.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, List, Optional, Tuple 2 | 3 | import numpy as np 4 | import torch 5 | import torch.fx 6 | from multipledispatch import dispatch 7 | 8 | from nnsmith.backends.factory import BackendCallable, BackendFactory 9 | from nnsmith.materialize.torch import TorchModel, numpify, get_cuda_string 10 | from nnsmith.materialize.torch.symbolnet import FxTracing 11 | 12 | 13 | class PT2(BackendFactory): 14 | def __init__(self, target: str = "cpu", optmax: bool = True, **kwargs): 15 | super().__init__(target, optmax) 16 | if self.target == "cpu": 17 | self.device = torch.device("cpu") 18 | elif self.target == "cuda": 19 | self.device = torch.device(get_cuda_string()) 20 | elif self.target == "mps": 21 | self.device = torch.device("mps") 22 | else: 23 | raise ValueError( 24 | f"Unknown target: {self.target}. Only `cpu` and `cuda` are supported." 25 | ) 26 | 27 | # get backend from kwargs or inductor by default 28 | self.backend = kwargs.get("backend", "inductor") 29 | self.mode = kwargs.get("mode") 30 | 31 | @property 32 | def system_name(self) -> str: 33 | return "pt2" 34 | 35 | @property 36 | def import_libs(self) -> List[str]: 37 | return ["import torch"] 38 | 39 | @dispatch(TorchModel) 40 | def make_backend(self, model: TorchModel) -> BackendCallable: 41 | torch_net = model.torch_model.to(self.device) 42 | # Names for parameters can be changed implicitly after compilation 43 | # We keep the original names to align with names in eager mode 44 | param_names = [k for k, _ in model.torch_model.named_parameters()] 45 | 46 | do_grad_check = model.needs_grad_check() 47 | 48 | if do_grad_check: 49 | torch_net = torch_net.train() 50 | else: 51 | torch_net = torch_net.eval() 52 | 53 | with torch.no_grad(): 54 | with FxTracing(): 55 | traced = torch.fx.symbolic_trace(torch_net) 56 | compiled = torch.compile( 57 | traced, fullgraph=True, backend=self.backend, mode=self.mode 58 | ) 59 | 60 | def closure(inputs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: 61 | nonlocal do_grad_check 62 | input_ts = [torch.from_numpy(v).to(self.device) for _, v in inputs.items()] 63 | if do_grad_check: 64 | outputs: List[torch.Tensor] = compiled(*input_ts) 65 | ret = {} 66 | 67 | for name, output in zip(torch_net.output_like.keys(), outputs): 68 | ret[name] = numpify(output) 69 | if output.requires_grad: 70 | # Get Vector-Jacobian product 71 | out_grad = torch.autograd.grad( 72 | outputs=output, 73 | inputs=compiled.parameters(), 74 | grad_outputs=torch.ones_like(output), 75 | retain_graph=True, 76 | allow_unused=True, 77 | ) 78 | for k, v in zip(param_names, out_grad): 79 | ret[name + "_vjp_" + k] = numpify(v) 80 | else: 81 | with torch.no_grad(): 82 | outputs: Tuple[torch.Tensor] = compiled(*input_ts) 83 | ret = {k: numpify(v) for k, v in zip(torch_net.output_like, outputs)} 84 | return ret 85 | 86 | return closure 87 | 88 | def emit_compile( 89 | self, opt_name: str, mod_name: str, inp_name: Optional[str] = None 90 | ) -> str: 91 | mode = f"'{self.mode}'" if self.mode else "None" 92 | return f"{opt_name} = torch.compile({mod_name}, fullgraph=True, backend='{self.backend}', mode={mode})" 93 | 94 | def emit_run(self, out_name: str, opt_name: str, inp_name: str) -> str: 95 | return f"""{out_name} = {opt_name}(*[torch.from_numpy(v).to('{self.device.type}') for v in {inp_name}]) 96 | {out_name} = [v.cpu().detach() for v in {out_name}] # torch2numpy 97 | {out_name} = [v.resolve_conj().numpy() if v.is_conj() else v.numpy() for v in {out_name}] # torch2numpy""" 98 | -------------------------------------------------------------------------------- /nnsmith/backends/torchjit.py: -------------------------------------------------------------------------------- 1 | import os 2 | import warnings 3 | from typing import Dict, List, Tuple 4 | 5 | import numpy as np 6 | import torch 7 | from multipledispatch import dispatch 8 | from torch.utils.mobile_optimizer import optimize_for_mobile 9 | 10 | from nnsmith.backends.factory import BackendCallable, BackendFactory 11 | from nnsmith.materialize.torch import TorchModel, numpify 12 | 13 | # Check https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html 14 | # for more PyTorch-internal options. 15 | NNSMITH_PTJIT_OPT_MOBILE = os.getenv("NNSMITH_PTJIT_OPT_MOBILE", "0") == "1" 16 | 17 | 18 | class TorchJIT(BackendFactory): 19 | def __init__(self, target="cpu", optmax: bool = False, **kwargs): 20 | super().__init__(target, optmax) 21 | if self.target == "cpu": 22 | self.device = torch.device("cpu") 23 | elif self.target == "cuda": 24 | self.device = torch.device("cuda") 25 | torch.backends.cudnn.benchmark = True 26 | else: 27 | raise ValueError(f"Unknown {target=}. Only `cpu` and `cuda` are supported.") 28 | 29 | @property 30 | def system_name(self) -> str: 31 | return "torchjit" 32 | 33 | @dispatch(TorchModel) 34 | def make_backend(self, model: TorchModel) -> BackendCallable: 35 | torch_net = model.torch_model.to(self.device) 36 | trace_inp = [ts.to(self.device) for ts in torch_net.get_random_inps().values()] 37 | 38 | do_grad_check = model.needs_grad_check() 39 | 40 | with warnings.catch_warnings(): 41 | warnings.simplefilter("ignore", category=torch.jit.TracerWarning) 42 | if do_grad_check: 43 | torch_net = torch_net.train() 44 | compiled = torch.jit.trace(torch_net, trace_inp) 45 | else: 46 | torch_net = torch_net.eval() 47 | with torch.no_grad(): 48 | compiled = torch.jit.trace(torch_net, trace_inp) 49 | compiled = torch.jit.freeze(compiled) # Frozen graph 50 | compiled = torch.jit.optimize_for_inference(compiled) 51 | 52 | if self.target == "cpu" and NNSMITH_PTJIT_OPT_MOBILE: 53 | compiled = optimize_for_mobile(compiled) 54 | 55 | def closure(inputs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: 56 | nonlocal do_grad_check 57 | input_ts = [torch.from_numpy(v).to(self.device) for _, v in inputs.items()] 58 | if do_grad_check: 59 | outputs: List[torch.Tensor] = compiled(*input_ts) 60 | params = {k: v for k, v in compiled.named_parameters()} 61 | ret = {} 62 | 63 | for name, output in zip(torch_net.output_like.keys(), outputs): 64 | ret[name] = numpify(output) 65 | if output.requires_grad: 66 | # get Vector-Jacobian product 67 | out_grad = torch.autograd.grad( 68 | outputs=output, 69 | inputs=params.values(), 70 | grad_outputs=torch.ones_like(output), 71 | retain_graph=True, 72 | allow_unused=True, 73 | ) 74 | for k, v in zip(params.keys(), out_grad): 75 | ret[name + "_vjp_" + k] = numpify(v) 76 | else: 77 | with torch.no_grad(): 78 | outputs: Tuple[torch.Tensor] = compiled(*input_ts) 79 | ret = {k: numpify(v) for k, v in zip(torch_net.output_like, outputs)} 80 | return ret 81 | 82 | return closure 83 | 84 | @property 85 | def import_libs(self) -> List[str]: 86 | return ["import torch"] 87 | 88 | def emit_compile(self, opt_name: str, mod_name: str, inp_name: str) -> str: 89 | return f"{opt_name} = torch.jit.trace({mod_name}, [torch.from_numpy(v).to('{self.device.type}') for v in {inp_name}])" 90 | 91 | def emit_run(self, out_name: str, opt_name: str, inp_name: str) -> str: 92 | return f"""{out_name} = {opt_name}(*[torch.from_numpy(v).to('{self.device.type}') for v in {inp_name}]) 93 | {out_name} = [v.cpu().detach() for v in {out_name}] # torch2numpy 94 | {out_name} = [v.resolve_conj().numpy() if v.is_conj() else v.numpy() for v in {out_name}] # torch2numpy""" 95 | -------------------------------------------------------------------------------- /doc/concept.md: -------------------------------------------------------------------------------- 1 | # Development Guide of NNSmith 2 | 3 | ## Bug Report Format 4 | 5 | *To be standardarized.* 6 | 7 | - **Bug-introducing model**: such as `model.onnx` for ONNX models and a saved-model folder in TensorFlow. 8 | - **Oracle**: `oracle.pkl` contains a dictionary 9 | - `"input"`: A dictionary of input. 10 | - `"output"`: A dictionary of output if the results are expected to be compared or `None` if the output contains `NaN` or `Inf` as undefined behaviours. 11 | - **Meta information**: `meta.json` meta information. 12 | - `"system"`: like `"tvm-gpu"` and `"ort-cpu"` 13 | - `"version"`: a version string hooked from `${SystemPackage}.__version__` 14 | - `"symptom"`: `"crash"` or `"inconsistency"` or `"timeout"` 15 | - `"version_id"` (optional): an identifier of the system's version (e.g., git hash or version strings) 16 | 17 | ## Abstract Operators (AO) 18 | 19 | An abstract operator contains information to construct a materialized operator. We will start with a simplified pooling 2D example. 20 | 21 | ### `__init__` 22 | 23 | The initializer of Abstract Operators (AO) takes a list of symbolic integers. In this way, during model generation, we can construct operators by feeding a certain number of symbolic integers. 24 | 25 | ```python 26 | class Pool2d(UnaryOpBase): 27 | def __init__(self, kw, kh, stride, pad): 28 | # Step 1: Invoke daddy's constructor 29 | super().__init__() 30 | 31 | # Step 2: Take arguments as attributes 32 | self.kw, self.kh = kw, kh 33 | self.stride = stride 34 | self.pad = pad 35 | 36 | # Step 3: Define desired operator input and output ranks 37 | # Typing: List[Tuple] where each tuple has some ranks 38 | # [ input.0(ok_rank.0, ok_rank.1, ...), input.1(...), ... ] 39 | # Why [(4,)]? 40 | # 1. Pooling2D only takes one input/output => one tuple; 41 | # 2. Pooling2D only accepts NCHW tensors => the only viable dim is 4; 42 | self.inp_ranks = [(4,)] 43 | self.out_ranks = [(4,)] 44 | ``` 45 | 46 | ### `type_transfer(itensors{.shape, .dtype}) => otensors{.shape, .dtype}` 47 | 48 | `type_transfer` is a type inference function to infer the output type (shape and data type) given inputs’ type information. 49 | 50 | ```python 51 | def type_infer(self, itensors: List[AbsTensor]): 52 | n, c, h, w = itensors[0].shape 53 | return [ # List 54 | AbsTensor(shape=[n, c, 55 | ((h - self.kh) + 2 * self.pad) // self.stride, 56 | ((w - self.kw) + 2 * self.pad) // self.stride, 57 | ], dtype=itensors[0].dtype)] 58 | ``` 59 | 60 | ### `requires(itensors{.shape, .dtype}) => [constraints, ...]` 61 | 62 | `requires` returns constraints (predicates) that must be satisfied when inserting this operator into a computational graph. 63 | 64 | ```python 65 | def requires(self, itensors: List[AbsTensor]): 66 | n, c, h, w = itensors[0].shape 67 | return [ # List 68 | self.kh >= 1, self.kw >= 1, self.stride >= 1, self.pad >= 0, 69 | self.kh <= h + 2 * self.pad, 70 | self.kw <= w + 2 * self.pad, 71 | ] 72 | ``` 73 | 74 | ### Class members 75 | 76 | - Viable data types: 77 | - `inp_dtypes`: Similar to `self.inp_ranks`, it contains a list of independent and viable input data types. 78 | - `out_dtypes` 79 | 80 | ### Varadic operator parameter 81 | 82 | Sometimes, we want to define an AO that can take variadic numbers of arguments, e.g., `Padding` can take a padding list which can be of 4 pad sizes (for (H and W) x (left-most and right-most)) for 4D tensors (e.g., NCHW) and 6 pad sizes for 5D tensors. 83 | 84 | Therefore, we need to let the model generator know how many arguments / symbolic integers a (padding) operator accepts. To specify this information, we overload the class data member `num_var_param: List[int]` which takes a list of integers, each of which is an acceptable number of arguments. For example, to create a padding operator that accepts 4 ~ 6D tensors. 85 | 86 | ### Constraining viable ranks of different input tensors 87 | 88 | Operators like `Concat` require input tensors to have the same shape (and thus ranks). 89 | 90 | To do so, we overload the `same_inp_dims` class variable to `True`: 91 | 92 | ```python 93 | class Concat(AbsOpBase) 94 | ... 95 | same_inp_dims = True 96 | ... 97 | def __init__(...): 98 | ... 99 | ``` 100 | -------------------------------------------------------------------------------- /nnsmith/backends/tvm.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from typing import List, Optional 3 | 4 | import tvm 5 | from multipledispatch import dispatch 6 | from tvm import relay 7 | 8 | from nnsmith.backends import BackendFactory 9 | from nnsmith.materialize.onnx import ONNXModel 10 | 11 | logging.getLogger("te_compiler").disabled = True 12 | logging.getLogger("autotvm").disabled = True 13 | 14 | 15 | def list_eq(a, b): 16 | if len(a) != len(b): 17 | return False 18 | for i in range(len(a)): 19 | if a[i] != b[i]: 20 | return False 21 | return True 22 | 23 | 24 | class TVM(BackendFactory): 25 | def __init__( 26 | self, 27 | target: str = "cpu", 28 | optmax: bool = True, 29 | executor: str = "graph", 30 | **kwargs, 31 | ) -> None: 32 | super().__init__(target, optmax, **kwargs) 33 | # WARNING: setting opt_level 4 sometimes causes false alarms 34 | # as in this level fast_math is enabled where slight numerical 35 | # inconsistency is allowed and outputs for UB-input may differ. 36 | self.opt_level = 4 if optmax else 0 37 | self.target = target 38 | if target == "cpu": 39 | self.tvm_target = tvm.target.Target("llvm") 40 | else: 41 | tvm_possible_targets = tvm.target.Target.list_kinds() 42 | assert ( 43 | target in tvm_possible_targets 44 | ), f"Unknown target {target}. Possible targets are {tvm_possible_targets}" 45 | self.tvm_target = tvm.target.Target(target) 46 | 47 | self.executor_mode = executor 48 | 49 | def get_device(self): 50 | dev_cand = self.tvm_target.export()["keys"] 51 | assert len(dev_cand) > 0, f"No viable device found for {self.tvm_target}" 52 | if "cuda" in dev_cand: 53 | return tvm.cuda() 54 | if "rocm" in dev_cand: 55 | return tvm.rocm() 56 | if "cpu" in dev_cand: 57 | return tvm.cpu() 58 | return tvm.device(dev_cand[0]) 59 | 60 | @property 61 | def system_name(self) -> str: 62 | return "tvm" 63 | 64 | @staticmethod 65 | def cvt_result(output): 66 | """Pack output tensor(s) into a list""" 67 | # TODO(jinkun): may not work for nested list / dynamic shape 68 | assert output is not None, "Output should not be None" 69 | if isinstance(output, (tvm.runtime.container.ADT, list)): 70 | output = [r.numpy() for r in output] 71 | elif output is not None: 72 | output = [output.numpy()] 73 | return output 74 | 75 | @dispatch(ONNXModel) 76 | def make_backend(self, model: ONNXModel): 77 | onnx_model = model.native_model 78 | shape_dict = {name: aten.shape for name, aten in model.input_like.items()} 79 | mod, params = relay.frontend.from_onnx( 80 | onnx_model, shape_dict, freeze_params=True 81 | ) 82 | mod = relay.transform.InferType()(mod) 83 | 84 | with tvm.transform.PassContext(opt_level=self.opt_level): 85 | executor = relay.build_module.create_executor( 86 | self.executor_mode, mod, self.get_device(), self.tvm_target, params 87 | ).evaluate() 88 | 89 | def closure(inputs): 90 | output = executor(**inputs) 91 | output = self.cvt_result(output) 92 | return dict(zip(model.output_like.keys(), output)) 93 | 94 | return closure 95 | 96 | @property 97 | def import_libs(self) -> List[str]: 98 | return ["import tvm", "from tvm import relay"] 99 | 100 | @property 101 | def version(self) -> str: 102 | return tvm.__version__ 103 | 104 | def emit_compile( 105 | self, opt_name: str, mod_name: str, inp_name: Optional[str] = None 106 | ) -> str: 107 | tab = " " * 4 108 | s = f"mod, params = relay.frontend.from_onnx({mod_name}, shape_dict, freeze_params=True)\n" 109 | s += f"with tvm.transform.PassContext(opt_level={self.opt_level}):\n" \ 110 | f"{tab}{opt_name} = relay.build_module.create_executor({self.executor_mode}, " \ 111 | f"mod, \'{self.get_device()}\', \'{self.target}\', params).evaluate()\n" 112 | return s 113 | 114 | def emit_run(self, out_name: str, opt_name: str, inp_name: str) -> str: 115 | return f"{out_name} = {opt_name}({inp_name})" 116 | -------------------------------------------------------------------------------- /nnsmith/materialize/torch/input_gen.py: -------------------------------------------------------------------------------- 1 | import time 2 | from abc import ABC, abstractmethod 3 | from typing import Dict, List, Tuple, Optional 4 | 5 | import numpy as np 6 | import torch 7 | 8 | from nnsmith.abstract.op import DType 9 | from nnsmith.materialize.torch.symbolnet import SymbolNet, random_tensor 10 | 11 | 12 | class InputSearchBase(ABC): 13 | @staticmethod 14 | def apply_weights(net, weight_sample): 15 | with torch.no_grad(): 16 | for name, param in net.named_parameters(): 17 | param.copy_(weight_sample[name]) 18 | 19 | def __init__( 20 | self, net: SymbolNet, start_inputs=None, start_weights=None, use_cuda: Optional[torch.device] = None 21 | ): 22 | self.net = net 23 | self.start_inputs = start_inputs 24 | self.start_weights = start_weights 25 | self.use_cuda = use_cuda 26 | 27 | @abstractmethod 28 | def search_one(self, start_inp, timeout_ms: int = None) -> Dict[str, torch.Tensor]: 29 | pass 30 | 31 | def search( 32 | self, max_time_ms: int = None, max_sample: int = 1 33 | ) -> Tuple[int, Dict[str, torch.Tensor]]: 34 | n_try = 0 35 | sat_inputs = None 36 | start_time = time.time() 37 | 38 | while ( 39 | max_time_ms is None or time.time() - start_time < max_time_ms / 1000 40 | ) and n_try < max_sample: 41 | if self.start_weights is not None and n_try < len(self.start_weights): 42 | self.apply_weights(self.net, self.start_weights[n_try]) 43 | else: 44 | weight_sample = {} 45 | for name, param in self.net.named_parameters(): 46 | weight_sample[name] = random_tensor( 47 | param.shape, dtype=param.dtype, use_cuda=self.use_cuda 48 | ) 49 | self.apply_weights(self.net, weight_sample) 50 | 51 | if self.start_inputs is not None and n_try < len(self.start_inputs): 52 | cur_input = self.start_inputs[n_try] 53 | else: 54 | cur_input = self.net.get_random_inps(use_cuda=self.use_cuda) 55 | 56 | res = self.search_one(cur_input, max_time_ms) 57 | n_try += 1 58 | if res is not None: 59 | sat_inputs = res 60 | break 61 | 62 | return n_try, sat_inputs 63 | 64 | 65 | class SamplingSearch(InputSearchBase): 66 | # Think about how people trivially generate inputs. 67 | def search_one(self, start_inp, timeout_ms: int = None) -> Dict[str, torch.Tensor]: 68 | with torch.no_grad(): 69 | self.net.check_intermediate_numeric = True 70 | _ = self.net(*start_inp.values()) 71 | if not self.net.invalid_found_last: 72 | return start_inp 73 | 74 | return None 75 | 76 | 77 | class GradSearch(InputSearchBase): 78 | def search_one(self, start_inp, timeout_ms: int = None) -> Dict[str, torch.Tensor]: 79 | timeout_s = None if timeout_ms is None else timeout_ms / 1000 80 | return self.net.grad_input_gen( 81 | init_tensors=start_inp, use_cuda=self.use_cuda, max_time=timeout_s 82 | ) 83 | 84 | 85 | class PracticalHybridSearch(InputSearchBase): 86 | def __init__( 87 | self, net: SymbolNet, start_inputs=None, start_weights=None, use_cuda: Optional[torch.device] = None 88 | ): 89 | super().__init__(net, start_inputs, start_weights, use_cuda) 90 | 91 | self.differentiable = None 92 | 93 | if all([DType.is_float(v.dtype) for _, v in self.net.input_like.items()]): 94 | diff_test_inp = self.net.get_random_inps(use_cuda=self.use_cuda) 95 | for _, item in diff_test_inp.items(): 96 | item.requires_grad_() 97 | self.net.forward(*diff_test_inp.values()) 98 | self.differentiable = self.net.differentiable 99 | else: 100 | self.differentiable = False 101 | 102 | def search_one(self, start_inp, timeout_ms: int = None) -> Dict[str, torch.Tensor]: 103 | # if this model is purely differentiable -> GradSearch 104 | # otherwise -> SamplingSearch 105 | # FIXME: Estimate gradient (e.g., proxy gradient) for non-differentiable inputs. 106 | 107 | if self.differentiable: 108 | return GradSearch.search_one(self, start_inp, timeout_ms) 109 | else: 110 | return SamplingSearch.search_one(self, start_inp, timeout_ms) 111 | -------------------------------------------------------------------------------- /nnsmith/cli/model_gen.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import random 4 | import time 5 | 6 | import hydra 7 | import jsonpickle 8 | from omegaconf import DictConfig 9 | 10 | from nnsmith.abstract.extension import activate_ext 11 | from nnsmith.backends.factory import BackendFactory 12 | from nnsmith.graph_gen import SymbolicGen, model_gen, viz 13 | from nnsmith.logging import MGEN_LOG 14 | from nnsmith.materialize import Model, TestCase 15 | from nnsmith.narrow_spec import auto_opset 16 | from nnsmith.util import hijack_patch_requires, mkdir, op_filter 17 | 18 | from equality_saturation_helper.helper import EquivalentGraphHelper, FFIHelper 19 | 20 | 21 | @hydra.main(version_base=None, config_path="../config", config_name="main") 22 | def main(cfg: DictConfig): 23 | for _ in range(1): 24 | # Generate a random ONNX model 25 | # TODO(@ganler): clean terminal outputs. 26 | mgen_cfg = cfg["mgen"] 27 | 28 | seed = random.getrandbits(32) if mgen_cfg["seed"] is None else mgen_cfg["seed"] 29 | 30 | MGEN_LOG.info(f"Using seed {seed}") 31 | 32 | # TODO(@ganler): skip operators outside of model gen with `cfg[exclude]` 33 | model_cfg = cfg["model"] 34 | ModelType = Model.init(model_cfg["type"], backend_target=cfg["backend"]["target"]) 35 | ModelType.add_seed_setter() 36 | 37 | if cfg["backend"]["type"] is not None: 38 | factory = BackendFactory.init( 39 | cfg["backend"]["type"], 40 | target=cfg["backend"]["target"], 41 | optmax=cfg["backend"]["optmax"], 42 | parse_name=True, 43 | ) 44 | else: 45 | factory = None 46 | 47 | # GENERATION 48 | opset = auto_opset( 49 | ModelType, 50 | factory, 51 | vulops=mgen_cfg["vulops"], 52 | grad=mgen_cfg["grad_check"], 53 | ) 54 | opset = op_filter(opset, mgen_cfg["include"], mgen_cfg["exclude"]) 55 | hijack_patch_requires(mgen_cfg["patch_requires"]) 56 | activate_ext(opset=opset, factory=factory) 57 | 58 | tgen_begin = time.time() 59 | 60 | gen = model_gen( 61 | opset=opset, 62 | method=mgen_cfg["method"], 63 | seed=seed, 64 | max_elem_per_tensor=mgen_cfg["max_elem_per_tensor"], 65 | max_nodes=mgen_cfg["max_nodes"], 66 | timeout_ms=mgen_cfg["timeout_ms"], 67 | rank_choices=mgen_cfg["rank_choices"], 68 | dtype_choices=mgen_cfg["dtype_choices"], 69 | ) 70 | tgen = time.time() - tgen_begin 71 | 72 | if isinstance(gen, SymbolicGen): 73 | MGEN_LOG.info( 74 | f"{len(gen.last_solution)} symbols and {len(gen.solver.assertions())} constraints." 75 | ) 76 | 77 | if MGEN_LOG.getEffectiveLevel() <= logging.DEBUG: 78 | MGEN_LOG.debug("solution:" + ", ".join(map(str, gen.last_solution))) 79 | 80 | # MATERIALIZATION 81 | tmat_begin = time.time() 82 | 83 | ir = gen.make_concrete() 84 | MGEN_LOG.info(ir.to_dot()) 85 | 86 | ffi_helper = FFIHelper() 87 | helper = EquivalentGraphHelper(ffi_helper, ir) 88 | helper.initialize_inner_graph() 89 | helper.test_helper_lib() 90 | ir2 = helper.output_inner_graph() 91 | 92 | # MGEN_LOG.info(ir.to_dot()) 93 | 94 | MGEN_LOG.info( 95 | f"Generated DNN has {ir.n_var()} variables and {ir.n_compute_inst()} operators." 96 | ) 97 | 98 | mkdir(mgen_cfg["save"], yes="y") 99 | if cfg["debug"]["viz"]: 100 | fmt = cfg["debug"]["viz_fmt"].replace(".", "") 101 | viz(ir, os.path.join(mgen_cfg["save"], f"graph.{fmt}")) 102 | 103 | model1 = ModelType.from_gir(ir) 104 | model1.refine_weights() # either random generated or gradient-based. 105 | model1.set_grad_check(mgen_cfg["grad_check"]) 106 | model2 = ModelType.from_gir(ir2) 107 | model2.refine_weights() 108 | model2.set_grad_check(mgen_cfg["grad_check"]) 109 | 110 | oracle = model2.make_oracle() 111 | tmat = time.time() - tmat_begin 112 | 113 | tsave_begin = time.time() 114 | testcase = TestCase(model1, oracle) 115 | testcase.dump(root_folder=mgen_cfg["save"]) 116 | tsave = time.time() - tsave_begin 117 | 118 | MGEN_LOG.info( 119 | f"Time: @Generation: {tgen:.2f}s @Materialization: {tmat:.2f}s @Save: {tsave:.2f}s" 120 | ) 121 | 122 | 123 | if __name__ == "__main__": 124 | main() 125 | -------------------------------------------------------------------------------- /nnsmith/materialize/torch/parse.py: -------------------------------------------------------------------------------- 1 | import operator 2 | from typing import Any, Dict, List, Union, cast 3 | 4 | import torch 5 | import torch._dynamo as dynamo 6 | import torch.fx as fx 7 | import torch.nn as nn 8 | import torch.utils._pytree as pytree 9 | from torch.fx.passes.shape_prop import ShapeProp 10 | 11 | from nnsmith.abstract.dtype import DType 12 | from nnsmith.abstract.op import ConcreteOp, Input 13 | from nnsmith.abstract.tensor import AbsTensor 14 | from nnsmith.gir import GraphIR, InstExpr 15 | 16 | 17 | class PropInterpreter(ShapeProp): 18 | def run_node(self, n: fx.node.Node) -> Any: 19 | result = super().run_node(n) 20 | n.meta["res"] = result 21 | return result 22 | 23 | 24 | def parse(model: nn.Module, *example_args: List[torch.Tensor]) -> GraphIR: 25 | gm: fx.GraphModule = dynamo.export(model, *example_args)[0] 26 | # store shape info on nodes 27 | sp = PropInterpreter(gm) 28 | sp.run(*example_args) 29 | 30 | def load_args(args: Union[List, Dict[str, Any]]) -> Union[List, Dict[str, Any]]: 31 | """ 32 | Map nodes to their outputs while keeping structures and other values the same. 33 | """ 34 | return torch.fx.graph.map_arg(args, lambda n: n.meta["res"]) 35 | 36 | named_modules = dict(gm.named_modules()) 37 | ir = GraphIR() 38 | name2retvals: Dict[str, List[str]] = {} 39 | for i_node, node in enumerate(gm.graph.nodes): 40 | node = cast(fx.node.Node, node) 41 | if node.op == "placeholder": 42 | iexpr = InstExpr(Input(dim=len(node.meta["res"].shape)), []) 43 | else: 44 | args_flatten, args_treespec = pytree.tree_flatten(node.args) 45 | kwargs_flatten, kwargs_treespec = pytree.tree_flatten(node.kwargs) 46 | input_nodes = [ 47 | a 48 | for a in (args_flatten + kwargs_flatten) 49 | if isinstance(a, fx.node.Node) 50 | ] 51 | input_valstrs = list(map(lambda n: name2retvals[n.name][0], input_nodes)) 52 | input_like = list( 53 | map( 54 | lambda ts: AbsTensor( 55 | shape=ts.shape, dtype=DType.from_torch(ts.dtype) 56 | ), 57 | pytree.tree_flatten( 58 | list(map(lambda n: n.meta["res"], input_nodes)) 59 | )[0], 60 | ) 61 | ) 62 | output_like = list( 63 | map( 64 | lambda ts: AbsTensor( 65 | shape=ts.shape, dtype=DType.from_torch(ts.dtype) 66 | ), 67 | pytree.tree_flatten(node.meta["res"])[0], 68 | ) 69 | ) 70 | nodes2empty = ( 71 | lambda n: ConcreteOp.empty if isinstance(n, fx.node.Node) else n 72 | ) 73 | args_wo_nodes = pytree.tree_map(nodes2empty, node.args) 74 | kwargs_wo_nodes = pytree.tree_map(nodes2empty, node.kwargs) 75 | if node.op == "call_function": 76 | if ( 77 | node.target is operator.getitem 78 | and isinstance(node.args[0], fx.node.Node) 79 | and not isinstance(node.args[0].meta["res"], torch.Tensor) 80 | ): 81 | name2retvals[node.name] = [ 82 | name2retvals[node.args[0].name][node.args[1]] 83 | ] 84 | continue 85 | else: 86 | target_str = node._pretty_print_target(node.target) 87 | elif node.op == "call_method": 88 | target_str = f"torch.Tensor.{node.target}" 89 | elif node.op == "call_module": 90 | target = named_modules[node.target] 91 | target_str = repr(target) 92 | if target.__module__.startswith("torch.nn.modules"): 93 | target_str = f"torch.nn.{target_str}" 94 | elif node.op == "get_attr": 95 | raise NotImplementedError(f"{node.op = }, {node.name = }") 96 | elif node.op == "output": 97 | continue 98 | else: 99 | raise ValueError(f"Unexpected {node.op = }") 100 | 101 | iexpr = InstExpr( 102 | ConcreteOp( 103 | target_str, args_wo_nodes, kwargs_wo_nodes, input_like, output_like 104 | ), 105 | input_valstrs, 106 | ) 107 | 108 | name2retvals[node.name] = ir.add_inst(iexpr).retvals() 109 | # end for 110 | return ir 111 | -------------------------------------------------------------------------------- /doc/cli.md: -------------------------------------------------------------------------------- 1 | ## Installation 2 | 3 | ```shell 4 | python3 -m pip install "nnsmith[torch,onnx,tvm,onnxruntime]" --upgrade 5 | # Or try the HEAD branch: 6 | # pip install "git+https://github.com/ise-uiuc/nnsmith@main#egg=nnsmith[torch,onnx,tvm,onnxruntime]" --upgrade 7 | ``` 8 | 9 | The core functionality of NNSmith is graph generation. 10 | Based on this, we provide powerful functionalities to do fuzzing and bug reporting. 11 | Therefore, the systems and model types you want to fuzz are installed as dependencies as is shown in labels inside "`[]`". 12 | 13 | We currently have model formats: 14 | - `torch` 15 | - `onnx` 16 | - `tensorflow` (experimental) 17 | 18 | and backends: 19 | - `tvm`: TVM 20 | - `onnxruntime`: ONNXRuntime 21 | - `trt`: TensorRT 22 | - `xla`: XLA 23 | - `tflite`: TFLite 24 | - `torchjit`: PyTorch JIT 25 | 26 | Meanwhile, the backend of `xla` and `tflite` is installed as part of TensorFlow. 27 | 28 | You can also have your own by extending `nnsmith.materialize.Model` and `nnsmith.backends.BackendFactory`. 29 | 30 | ## Graph generation 31 | 32 | ```shell 33 | # Generate 5-node onnx model. 34 | nnsmith.model_gen mgen.max_nodes=5 model.type=onnx debug.viz=true 35 | # See: nnsmith_output/* (default output folder) 36 | 37 | # TensorFlow model. 38 | nnsmith.model_gen debug.viz=true model.type=tensorflow 39 | 40 | # User-spec. output directory 41 | nnsmith.model_gen debug.viz=true model.type=tensorflow mgen.save=tf_output 42 | ``` 43 | 44 | ## Locally debug a model 45 | 46 | ```python 47 | # Generate a onnx model 48 | nnsmith.model_gen model.type=onnx mgen.max_nodes=5 49 | 50 | # Check the model 51 | pip install onnxruntime # use ONNXRuntime to execute the model 52 | nnsmith.model_exec model.type=onnx backend.type=onnxruntime model.path=nnsmith_output/model.onnx 53 | # `model.path` should point to the exact model, instead of a folder. 54 | # It will first compile and run to see if there's any bug. 55 | # By default it will search `oracle.pkl` and do verification. 56 | 57 | # Check the model and do diff testing with tvm 58 | nnsmith.model_exec model.type=onnx \ 59 | backend.type=onnxruntime \ 60 | model.path=nnsmith_output/model.onnx \ 61 | cmp.with='{type:tvm, optmax:true, target:cpu}' 62 | ``` 63 | 64 | ## Experimental: Gradient checking 65 | 66 | For `pt2` and `torchjit`, we have initial supports for examining the gradients. 67 | 68 | To enable that, just need to append `mgen.grad_check=true` to the examples illustrated above. 69 | 70 | ## Data type testing 71 | 72 | Many compilers do not support a full set of operators (in ONNX and TensorFlow). Thus, we infer the support set by doing single operator testing. 73 | 74 | ```shell 75 | # Infer the support set of onnxruntime to ONNX format. 76 | nnsmith.dtype_test model.type="onnx" backend.type="onnxruntime" 77 | # Results are often cached in `~/.cache/nnsmith`. 78 | ``` 79 | 80 | ## Fuzzing 81 | 82 | ```shell 83 | nnsmith.fuzz fuzz.time=30s model.type=onnx backend.type=tvm fuzz.root=fuzz_report debug.viz=true 84 | # Bug reports are stored in `./fuzz_report`. 85 | ``` 86 | 87 | ## Limit operator types, ranks and data types 88 | 89 | To limit: 90 | - rank only to be 4 (needed by Conv2d); 91 | - data type only to be float32; 92 | - only include Conv2d and ReLU. 93 | 94 | ```shell 95 | yes | nnsmith.model_gen model.type=torch mgen.method=symbolic-cinit \ 96 | mgen.rank_choices="[4]" \ 97 | mgen.dtype_choices="[f32]" \ 98 | mgen.include="[core.NCHWConv2d, core.ReLU]" \ 99 | debug.viz=true 100 | ``` 101 | 102 | ## Add extra constraints 103 | 104 | ```shell 105 | # Create patch file as `patch.py` 106 | echo 'from nnsmith.abstract.arith import nnsmith_lt 107 | from nnsmith.abstract.extension import patch_requires 108 | 109 | 110 | @patch_requires("global", "core.NCHWConv2d") 111 | def limit_conv2d(self, _): 112 | # let the kernels to be > 3 113 | return [nnsmith_lt(3, self.kernel_h_size), nnsmith_lt(3, self.kernel_w_size)] 114 | ' > patch.py 115 | # Apply the patch with `mgen.patch_requires=./tests/mock/requires_patch.py` (can also be a list of paths) 116 | yes | nnsmith.model_gen model.type=torch mgen.method=symbolic-cinit \ 117 | mgen.rank_choices="[4]" \ 118 | mgen.dtype_choices="[f32]" \ 119 | mgen.include="[core.NCHWConv2d, core.ReLU]" \ 120 | mgen.patch_requires=./patch.py \ 121 | debug.viz=true 122 | ``` 123 | 124 | ## Misc 125 | 126 | TensorFlow logging can be very noisy. Use `TF_CPP_MIN_LOG_LEVEL=3` as environmental variable to depress that. 127 | -------------------------------------------------------------------------------- /experiments/legacy/plot_inp_search.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import argparse 4 | 5 | import matplotlib.pyplot as plt 6 | import pandas as pd 7 | import numpy as np 8 | 9 | 10 | SMALL_SIZE = 10 11 | MEDIUM_SIZE = 15 12 | BIGGER_SIZE = 18 13 | 14 | plt.rc("font", size=SMALL_SIZE) # controls default text sizes 15 | plt.rc("axes", titlesize=MEDIUM_SIZE) # fontsize of the axes title 16 | plt.rc("axes", labelsize=MEDIUM_SIZE) # fontsize of the x and y labels 17 | plt.rc("xtick", labelsize=SMALL_SIZE) # fontsize of the tick labels 18 | plt.rc("ytick", labelsize=SMALL_SIZE) # fontsize of the tick labels 19 | plt.rc("legend", fontsize=MEDIUM_SIZE) # legend fontsize 20 | plt.rc("figure", titlesize=BIGGER_SIZE) # fontsize of the figure title 21 | 22 | 23 | if __name__ == "__main__": 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument("--root", type=str, required=True) 26 | parser.add_argument("--output", type=str, default="results") 27 | args = parser.parse_args() 28 | 29 | REGEX_PATTERN = "(\d+)-model-(\d+)-node-exp" 30 | res = re.match(REGEX_PATTERN, args.root) 31 | n_model, n_nodes = res.groups() 32 | 33 | # Plot data 34 | # X: Time 35 | # Y: Succ. Rate 36 | 37 | sampling_time = [] 38 | sampling_succ_rate = [] 39 | 40 | grad_time = [] 41 | grad_succ_rate = [] 42 | 43 | proxy_time = [] 44 | proxy_succ_rate = [] 45 | 46 | for f in os.listdir(args.root): 47 | if f.endswith(".csv") and f != "model_info.csv": 48 | data = pd.read_csv(os.path.join(args.root, f)) 49 | 50 | sampling_time.append(data["sampling-time"].mean()) 51 | sampling_succ_rate.append(data["sampling-succ"].mean()) 52 | grad_time.append(data["grad-time"].mean()) 53 | grad_succ_rate.append(data["grad-succ"].mean()) 54 | proxy_time.append(data["proxy-time"].mean()) 55 | proxy_succ_rate.append(data["proxy-succ"].mean()) 56 | elif f == "model_info.csv": 57 | data = pd.read_csv(os.path.join(args.root, f)) 58 | gentime = data["gen_time"] 59 | print(f"{gentime.mean()=}, {gentime.min()=}, {gentime.max()=}") 60 | 61 | def sort_by_time(time, succ_rate): 62 | time = np.array(time) 63 | succ_rate = np.array(succ_rate) 64 | sort_idx = time.argsort() 65 | return time[sort_idx], succ_rate[sort_idx] 66 | 67 | # sort succ rate by time 68 | sampling_time, sampling_succ_rate = sort_by_time(sampling_time, sampling_succ_rate) 69 | grad_time, grad_succ_rate = sort_by_time(grad_time, grad_succ_rate) 70 | proxy_time, proxy_succ_rate = sort_by_time(proxy_time, proxy_succ_rate) 71 | 72 | fig, (ax1, ax2) = plt.subplots( 73 | 2, 1, figsize=(8, 4), sharex=True, gridspec_kw={"height_ratios": [2.5, 1]} 74 | ) 75 | fig.subplots_adjust(hspace=0.1) 76 | 77 | for ax in [ax1, ax2]: 78 | ax.plot( 79 | sampling_time * 1000, 80 | sampling_succ_rate, 81 | markersize=10, 82 | marker="x", 83 | linestyle="--", 84 | label="Sampling", 85 | ) 86 | ax.plot( 87 | grad_time * 1000, 88 | grad_succ_rate, 89 | marker="*", 90 | markersize=10, 91 | linestyle="--", 92 | label="Gradient", 93 | ) 94 | ax.plot( 95 | proxy_time * 1000, 96 | proxy_succ_rate, 97 | marker="^", 98 | markersize=10, 99 | linestyle="--", 100 | label="Gradient (Proxy Deriv.)", 101 | ) 102 | ax.grid(True, linestyle="--", linewidth=0.5) 103 | 104 | ax1.set_ylim(0.66, 1) 105 | ax2.set_ylim(0.29, 0.43) 106 | 107 | ax1.spines.bottom.set_visible(False) 108 | ax2.spines.top.set_visible(False) 109 | ax1.xaxis.tick_top() 110 | ax1.tick_params(labeltop=False) # don't put tick labels at the top 111 | ax2.xaxis.tick_bottom() 112 | 113 | d = 0.5 # proportion of vertical to horizontal extent of the slanted line 114 | kwargs = dict( 115 | marker=[(-1, -d), (1, d)], 116 | markersize=12, 117 | linestyle="none", 118 | color="k", 119 | mec="k", 120 | mew=1, 121 | clip_on=False, 122 | ) 123 | ax1.plot([0, 1], [0, 0], transform=ax1.transAxes, **kwargs) 124 | ax2.plot([0, 1], [1, 1], transform=ax2.transAxes, **kwargs) 125 | 126 | plt.legend(loc="lower right") 127 | ax1.set_yticks([0.7, 0.8, 0.9, 1]) 128 | ax2.set_yticks([0.3, 0.4]) 129 | 130 | ax2.set_xlabel("Avg. Searching Time (millisecond)", fontweight="bold") 131 | ax2.xaxis.set_label_coords(0.5, 0.05, transform=fig.transFigure) 132 | ax2.set_ylabel("Success Rate", fontweight="bold") 133 | ax2.yaxis.set_label_coords(0.05, 0.5, transform=fig.transFigure) 134 | 135 | plt.savefig(os.path.join(args.output, f"input-search{n_model}-{n_nodes}.pdf")) 136 | plt.savefig(os.path.join(args.output, f"input-search{n_model}-{n_nodes}.png")) 137 | -------------------------------------------------------------------------------- /nnsmith/cli/test_eq_helper.py: -------------------------------------------------------------------------------- 1 | from nnsmith.cli.equivalent_fuzz import EquivalentFuzzingLoop 2 | from nnsmith.materialize import TestCase 3 | from nnsmith.logging import FUZZ_LOG 4 | from omegaconf import OmegaConf 5 | import time 6 | from equality_saturation_helper.helper import FFIHelper, EquivalentGraphHelper 7 | import random 8 | from enum import Enum 9 | 10 | 11 | class RunType(Enum): 12 | OK = 1, 13 | MAKE_TESTCASE_FAIL = 2, 14 | RUN_TESTCASE_FAIL = 3, 15 | 16 | 17 | class TestCorrectnessOfEQHelper(EquivalentFuzzingLoop): 18 | def run(self): 19 | print(self.opset) 20 | start_time = time.time() 21 | ffi_helper = FFIHelper() 22 | assert ffi_helper.is_helper_loaded() 23 | iteration = 0 24 | passed = 0 25 | bug = 0 26 | while time.time() - start_time < self.timeout_s: 27 | print(f"pass/bug/iteration: {passed}/{bug}/{iteration}") 28 | iteration += 1 29 | 30 | pass_flag = 0 31 | bug_flag = 0 32 | 33 | seed = random.getrandbits(32) 34 | ir = self.make_gir(seed=seed) 35 | eq_graph_helper = EquivalentGraphHelper(ffi_helper=ffi_helper, gir=ir) 36 | eq_graph_helper.initialize_inner_graph() 37 | eq_graph_helper.run_saturation() 38 | eq_graph_helper.test_helper_lib() 39 | 40 | ir2, graph_output_map = eq_graph_helper.randomly_generate_an_equivalent_graph() 41 | run_res = self.cmp_result(ir, ir2, graph_output_map) 42 | if run_res == RunType.OK or run_res == RunType.MAKE_TESTCASE_FAIL: 43 | pass_flag += 1 44 | elif run_res == RunType.RUN_TESTCASE_FAIL: 45 | bug_flag += 1 46 | FUZZ_LOG.warning(f"Failed model seed: {seed}") 47 | 48 | ir3, graph_output_map = eq_graph_helper.find_the_most_complex_equivalent_graph() 49 | run_res = self.cmp_result(ir, ir3, graph_output_map) 50 | if run_res == RunType.OK or run_res == RunType.MAKE_TESTCASE_FAIL: 51 | pass_flag += 1 52 | elif run_res == RunType.RUN_TESTCASE_FAIL: 53 | bug_flag += 1 54 | FUZZ_LOG.warning(f"Failed model seed: {seed}") 55 | 56 | ir4, graph_output_map = eq_graph_helper.find_the_most_simplified_equivalent_graph() 57 | run_res = self.cmp_result(ir, ir4, graph_output_map) 58 | if run_res == RunType.OK or run_res == RunType.MAKE_TESTCASE_FAIL: 59 | pass_flag += 1 60 | elif run_res == RunType.RUN_TESTCASE_FAIL: 61 | bug_flag += 1 62 | FUZZ_LOG.warning(f"Failed model seed: {seed}") 63 | 64 | if pass_flag == 3: 65 | passed += 1 66 | 67 | if bug_flag > 0: 68 | bug += 1 69 | 70 | def cmp_result(self, ir1, ir2, graph_output_map) -> RunType: 71 | try: 72 | testcase_1 = self.gir_to_testcase(ir1) 73 | testcase_2 = self.gir_to_testcase(ir2) 74 | except BaseException as e: 75 | print(e) 76 | return RunType.MAKE_TESTCASE_FAIL 77 | self.make_inputs_and_weights_consist(testcase_1, testcase_2) 78 | testcase_1 = TestCase(testcase_1.model, testcase_1.model.make_oracle(testcase_1.oracle.input)) 79 | testcase_2 = TestCase(testcase_2.model, testcase_2.model.make_oracle(testcase_2.oracle.input)) 80 | try: 81 | self.make_output_consist(testcase_1, testcase_2, graph_output_map) 82 | except BaseException as e: 83 | print(testcase_1.model.gir.to_dot()) 84 | print(testcase_2.model.gir.to_dot()) 85 | return RunType.MAKE_TESTCASE_FAIL 86 | res = self.validate_results(testcase_1, testcase_2, graph_output_map) 87 | if res is not None: 88 | print(res) 89 | return RunType.RUN_TESTCASE_FAIL 90 | else: 91 | return RunType.OK 92 | 93 | def yield_manual_test(self): 94 | pass 95 | 96 | 97 | def main(): 98 | cfg = {'model': {'type': 'torch', 'path': '???'}, 99 | 'mgen': {'max_nodes': 10, 'timeout_ms': 10000, 'vulops': False, 'method': 'symbolic-cinit', 100 | 'save': 'nnsmith_output', 'seed': None, 'max_elem_per_tensor': 65536, 'rank_choices': [2, 4], 101 | 'dtype_choices': None, 'include': None, 'exclude': ["Round", "Cast"], 'patch_requires': [], 102 | 'grad_check': False}, 103 | 'backend': {'type': 'torchjit', 'optmax': True, 'target': 'cpu'}, 'ad': {'type': None}, 104 | 'cache': {'topset': True}, 105 | 'debug': {'viz': True, 'viz_fmt': 'png'}, 106 | 'fuzz': {'time': '12h', 'root': 'eq_test_output', 'seed': None, 'crash_safe': False, 'test_timeout': None, 107 | 'save_test': None}, 'filter': {'type': [], 'patch': []}, 108 | 'cmp': {'equal_nan': True, 'raw_input': None, 'oracle': None, 109 | 'with': {'type': None, 'optmax': True, 'target': 'cpu'}, 'seed': None, 'bug_presence': 'report', 110 | 'save': None}} 111 | cfg = OmegaConf.create(cfg) 112 | test = TestCorrectnessOfEQHelper(cfg) 113 | test.run() 114 | 115 | 116 | if __name__ == "__main__": 117 | main() 118 | -------------------------------------------------------------------------------- /nnsmith/util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import shutil 4 | from importlib.util import module_from_spec, spec_from_file_location 5 | from os import PathLike 6 | from typing import Callable, Dict, List, Union 7 | 8 | import numpy as np 9 | from omegaconf import ListConfig 10 | 11 | from nnsmith.logging import MGEN_LOG 12 | 13 | try: 14 | import pygraphviz as pgv 15 | 16 | HAS_PYGRAPHVIZ = True 17 | except ImportError: 18 | import warnings 19 | 20 | warnings.warn( 21 | "Install pygraphviz for visualization: https://pygraphviz.github.io/documentation/stable/install.html\n" 22 | "Currently graph visualization is not enabled." 23 | ) 24 | pgv = None 25 | HAS_PYGRAPHVIZ = False 26 | 27 | 28 | from nnsmith.logging import CORE_LOG, VIZ_LOG 29 | 30 | SEED_SETTERS = { 31 | "random": random.seed, 32 | "numpy": np.random.seed, 33 | } 34 | 35 | 36 | def register_seed_setter( 37 | name: str, 38 | fn: Callable[[int], None], 39 | overwrite=False, 40 | ): 41 | if not overwrite: 42 | assert name not in SEED_SETTERS, f"{name} is already registered" 43 | SEED_SETTERS[name] = fn 44 | CORE_LOG.debug(f"Register seed setter for {name}") 45 | 46 | 47 | def set_seed(seed: int, names: List = None): 48 | if names is None: 49 | names = SEED_SETTERS.keys() 50 | for name in names: 51 | SEED_SETTERS[name](seed) 52 | 53 | 54 | def mkdir(dir: os.PathLike, yes=False): 55 | if os.path.exists(dir): 56 | decision = "" 57 | if yes: 58 | decision = "y" 59 | while decision.lower() not in ["y", "n"]: 60 | CORE_LOG.warning( 61 | "Report folder already exists. Press [Y/N] to continue or exit..." 62 | ) 63 | decision = input() 64 | if decision.lower() == "n": 65 | CORE_LOG.error(f"{dir} already exist... Remove it or use a different name.") 66 | raise RuntimeError("Folder already exists") 67 | else: 68 | shutil.rmtree(dir) 69 | 70 | os.makedirs(dir) 71 | 72 | 73 | def parse_timestr(timestr: str): 74 | if timestr.endswith("hr"): 75 | return int(timestr[:-2]) * 3600 76 | elif timestr.endswith("h"): 77 | return int(timestr[:-1]) * 3600 78 | elif timestr.endswith("min"): 79 | return int(timestr[:-3]) * 60 80 | elif timestr.endswith("m"): 81 | return int(timestr[:-1]) * 60 82 | elif timestr.endswith("s"): 83 | return int(timestr[:-1]) 84 | 85 | raise ValueError( 86 | f"Cannot parse time string: {timestr}. Valid examples: 1hr, 1h, 1min, 1m, 1s" 87 | ) 88 | 89 | 90 | def is_invalid(output: Dict[str, np.ndarray]): 91 | for _, o in output.items(): 92 | if np.isnan(o).any() or np.isinf(o).any(): 93 | return True 94 | return False 95 | 96 | 97 | _DOT_EXIST = shutil.which("dot") is not None 98 | _CONDA_EXIST = shutil.which("conda") is not None 99 | _APT_EXIST = shutil.which("apt") is not None 100 | _BREW_EXIST = shutil.which("brew") is not None 101 | 102 | _CALL_ONCE = False 103 | 104 | 105 | def _check_dot_install(): 106 | global _CALL_ONCE 107 | if not _DOT_EXIST and not _CALL_ONCE: 108 | _CALL_ONCE = True 109 | VIZ_LOG.warning("`dot` not found.") 110 | if _CONDA_EXIST or _APT_EXIST or _BREW_EXIST: 111 | VIZ_LOG.warning("To install via:") 112 | if _CONDA_EXIST: 113 | VIZ_LOG.warning(" conda:\t conda install -c anaconda graphviz -y") 114 | 115 | if _APT_EXIST: 116 | VIZ_LOG.warning(" apt:\t sudo apt install graphviz -y") 117 | 118 | if _BREW_EXIST: 119 | VIZ_LOG.warning(" brew:\t brew install graphviz") 120 | 121 | VIZ_LOG.warning("Also see: https://graphviz.org/download/") 122 | return False 123 | 124 | return True 125 | 126 | 127 | def viz_dot(dotobj, filename: str = None): 128 | if _check_dot_install(): 129 | if filename is None: 130 | filename = f"graph.png" 131 | 132 | if isinstance(dotobj, str): 133 | dotobj = pgv.AGraph(dotobj) 134 | 135 | dotobj.layout("dot") 136 | dotobj.draw(filename) 137 | 138 | 139 | def op_filter(topset, include=None, exclude=None): 140 | if include is not None and exclude is not None: 141 | # use either include or exclude 142 | raise ValueError("Cannot use both include and exclude") 143 | 144 | if include: 145 | return [op for op in topset if op.name() in include] 146 | 147 | if exclude: 148 | return [op for op in topset if op.name() not in exclude] 149 | 150 | return topset 151 | 152 | 153 | def hijack_patch_requires(patch_paths: Union[PathLike, List[PathLike]]): 154 | patch_paths = ( 155 | patch_paths if isinstance(patch_paths, (ListConfig, list)) else [patch_paths] 156 | ) 157 | for f in patch_paths: 158 | assert os.path.isfile( 159 | f 160 | ), "mgen.requires_patch must be a list of file locations." 161 | text = open(f).read() 162 | assert ( 163 | "@patch_requires(" in text 164 | ), f"No patch found in the {f} after checking `@patch_requires(`" 165 | spec = spec_from_file_location("nnsmith.ext.patch_requires", f) 166 | spec.loader.exec_module(module_from_spec(spec)) 167 | MGEN_LOG.info(f"Import patch_requires from {f}") 168 | -------------------------------------------------------------------------------- /experiments/legacy/nnsmith_gen_onnx.py: -------------------------------------------------------------------------------- 1 | from nnsmith.graph_gen import random_model_gen, SymbolNet 2 | from nnsmith.materialize.torch.input_gen import PracticalHybridSearch 3 | from nnsmith.materialize.onnx.export import torch2onnx 4 | from nnsmith.dtype_test import rewrite_op_dtype 5 | from nnsmith.abstract.op import ALL_OP_TYPES 6 | from nnsmith.util import mkdir 7 | 8 | from experiments.graphfuzz import GraphFuzz 9 | 10 | import pickle 11 | import os 12 | import random 13 | import argparse 14 | import time 15 | import warnings 16 | import tarfile 17 | 18 | from tqdm import tqdm 19 | import torch 20 | 21 | 22 | def nnsmith_gen_once( 23 | path_prefix, seed, max_nodes, candidates_overwrite=None, mode="random" 24 | ): 25 | if mode == "hybrid": 26 | mode = random.choice(["random", "guided"]) 27 | 28 | torch.manual_seed(seed) 29 | gen_tstart = time.time() 30 | gen, solution = random_model_gen( 31 | init_rank=4, 32 | seed=seed, 33 | max_nodes=max_nodes, 34 | candidates_overwrite=candidates_overwrite, 35 | mode=mode, 36 | ) 37 | net = SymbolNet( 38 | gen.abstract_graph, solution, verbose=False, alive_shapes=gen.alive_shapes 39 | ) 40 | gen_time = time.time() - gen_tstart 41 | 42 | net.enable_proxy_grad() 43 | net.eval() # otherwise BN wants batch > 1 44 | searcher = PracticalHybridSearch(net) 45 | n_try, sat_inputs = searcher.search( 46 | max_time_ms=gen_time * 0.02 * 1000, max_sample=2, return_list=True 47 | ) 48 | net.disable_proxy_grad() 49 | 50 | with torch.no_grad(): 51 | net.eval() 52 | 53 | test_inputs = sat_inputs if sat_inputs else net.get_random_inps(use_cuda=False) 54 | 55 | outputs = net.forward(*test_inputs) 56 | 57 | inames, onames, oidx = torch2onnx( 58 | net, 59 | path_prefix + ".onnx", 60 | verbose=False, 61 | use_cuda=False, 62 | dummy_inputs=test_inputs, 63 | ) 64 | 65 | inputs = [t.cpu().numpy() for t in test_inputs] 66 | 67 | if isinstance(outputs, torch.Tensor): 68 | outputs = [outputs.cpu().numpy()] 69 | else: 70 | outputs = [o.cpu().numpy() for o in outputs] 71 | 72 | input_dict = {ina: inp for ina, inp in zip(inames, inputs)} 73 | output_dict = {oname: outputs[i] for oname, i in zip(onames, oidx)} 74 | 75 | with open(path_prefix + ".pkl", "wb") as f: 76 | pickle.dump((input_dict, output_dict), f) 77 | 78 | 79 | if __name__ == "__main__": 80 | parser = argparse.ArgumentParser() 81 | parser.add_argument("--onnx_dir", type=str, required=True) 82 | parser.add_argument("--time_budget", type=int, default=60 * 60 * 4) 83 | parser.add_argument("--max_nodes", type=int, default=10) 84 | parser.add_argument("--graphfuzz_ops", action="store_true") 85 | parser.add_argument("--ort_cache", type=str, default=None) 86 | parser.add_argument("--seed", type=int, default=233) 87 | parser.add_argument("--mode", type=str, default="random") 88 | parser.add_argument("--tar", action="store_true") 89 | args = parser.parse_args() 90 | 91 | mkdir(args.onnx_dir) 92 | 93 | random.seed(args.seed) 94 | torch.manual_seed(args.seed) 95 | 96 | if args.ort_cache: 97 | print(args.ort_cache) 98 | if not os.path.exists(args.ort_cache): 99 | print(f"Please first generate cache! (mkdir config first)") 100 | print(f"python nnsmith/dtype_test.py --cache {args.ort_cache}") 101 | exit(1) 102 | # must pre run this. otherwise using ort will slow down generation. 103 | rewrite_op_dtype(ALL_OP_TYPES, factory=None, cache=args.ort_cache) 104 | 105 | if args.graphfuzz_ops: 106 | candidates_overwrite = GraphFuzz.get_available_op_ts() 107 | else: 108 | candidates_overwrite = None 109 | 110 | # FORMAT: {generation time cost in seconds}, {model relative path} 111 | # MUST RANK by GENERATION ORDER. 112 | config_file = open(os.path.join(args.onnx_dir, "gentime.csv"), "w") 113 | 114 | start_time = time.time() 115 | gen_cnt = 0 116 | valid_cnt = 0 117 | 118 | if args.tar: 119 | tar = tarfile.open(os.path.join(args.onnx_dir, "models.tar"), "w") 120 | 121 | with tqdm(total=args.time_budget) as pbar: 122 | while time.time() - start_time < args.time_budget: 123 | seed = random.getrandbits(32) 124 | 125 | tstart = time.time() 126 | try: 127 | with warnings.catch_warnings(): # just shutup. 128 | warnings.simplefilter("ignore") 129 | nnsmith_gen_once( 130 | os.path.join(args.onnx_dir, f"{valid_cnt}"), 131 | seed, 132 | max_nodes=10, 133 | candidates_overwrite=candidates_overwrite, 134 | mode=args.mode, 135 | ) 136 | to_name = f"{valid_cnt}.onnx" 137 | label = to_name 138 | valid_cnt += 1 139 | if args.tar: 140 | tar.add( 141 | os.path.join(args.onnx_dir, to_name), 142 | arcname="models/" + to_name, 143 | ) 144 | os.unlink(os.path.join(args.onnx_dir, to_name)) 145 | except Exception as e: 146 | print(f"Fail when seed={seed}") 147 | print(e) # Skip a few errors. 148 | label = "FAILURE" 149 | 150 | time_diff = time.time() - tstart 151 | config_file.write(f"{time_diff:.5f},{label}\n") 152 | 153 | gen_cnt += 1 154 | config_file.flush() 155 | 156 | pbar.update(int(time.time() - start_time) - pbar.n) 157 | pbar.set_description(f"valid={valid_cnt},fail={gen_cnt-valid_cnt}") 158 | pbar.refresh() 159 | config_file.close() 160 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | 2 | # Contributor Covenant Code of Conduct 3 | 4 | ## Our Pledge 5 | 6 | We as members, contributors, and leaders pledge to make participation in our 7 | community a harassment-free experience for everyone, regardless of age, body 8 | size, visible or invisible disability, ethnicity, sex characteristics, gender 9 | identity and expression, level of experience, education, socio-economic status, 10 | nationality, personal appearance, race, caste, color, religion, or sexual 11 | identity and orientation. 12 | 13 | We pledge to act and interact in ways that contribute to an open, welcoming, 14 | diverse, inclusive, and healthy community. 15 | 16 | ## Our Standards 17 | 18 | Examples of behavior that contributes to a positive environment for our 19 | community include: 20 | 21 | * Demonstrating empathy and kindness toward other people 22 | * Being respectful of differing opinions, viewpoints, and experiences 23 | * Giving and gracefully accepting constructive feedback 24 | * Accepting responsibility and apologizing to those affected by our mistakes, 25 | and learning from the experience 26 | * Focusing on what is best not just for us as individuals, but for the overall 27 | community 28 | 29 | Examples of unacceptable behavior include: 30 | 31 | * The use of sexualized language or imagery, and sexual attention or advances of 32 | any kind 33 | * Trolling, insulting or derogatory comments, and personal or political attacks 34 | * Public or private harassment 35 | * Publishing others' private information, such as a physical or email address, 36 | without their explicit permission 37 | * Other conduct which could reasonably be considered inappropriate in a 38 | professional setting 39 | 40 | ## Enforcement Responsibilities 41 | 42 | Community leaders are responsible for clarifying and enforcing our standards of 43 | acceptable behavior and will take appropriate and fair corrective action in 44 | response to any behavior that they deem inappropriate, threatening, offensive, 45 | or harmful. 46 | 47 | Community leaders have the right and responsibility to remove, edit, or reject 48 | comments, commits, code, wiki edits, issues, and other contributions that are 49 | not aligned to this Code of Conduct, and will communicate reasons for moderation 50 | decisions when appropriate. 51 | 52 | ## Scope 53 | 54 | This Code of Conduct applies within all community spaces, and also applies when 55 | an individual is officially representing the community in public spaces. 56 | Examples of representing our community include using an official e-mail address, 57 | posting via an official social media account, or acting as an appointed 58 | representative at an online or offline event. 59 | 60 | ## Enforcement 61 | 62 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 63 | reported to the community leaders responsible for enforcement at 64 | jiawei6@illinois.edu / jaway.liu@gmail.com. 65 | All complaints will be reviewed and investigated promptly and fairly. 66 | 67 | All community leaders are obligated to respect the privacy and security of the 68 | reporter of any incident. 69 | 70 | ## Enforcement Guidelines 71 | 72 | Community leaders will follow these Community Impact Guidelines in determining 73 | the consequences for any action they deem in violation of this Code of Conduct: 74 | 75 | ### 1. Correction 76 | 77 | **Community Impact**: Use of inappropriate language or other behavior deemed 78 | unprofessional or unwelcome in the community. 79 | 80 | **Consequence**: A private, written warning from community leaders, providing 81 | clarity around the nature of the violation and an explanation of why the 82 | behavior was inappropriate. A public apology may be requested. 83 | 84 | ### 2. Warning 85 | 86 | **Community Impact**: A violation through a single incident or series of 87 | actions. 88 | 89 | **Consequence**: A warning with consequences for continued behavior. No 90 | interaction with the people involved, including unsolicited interaction with 91 | those enforcing the Code of Conduct, for a specified period of time. This 92 | includes avoiding interactions in community spaces as well as external channels 93 | like social media. Violating these terms may lead to a temporary or permanent 94 | ban. 95 | 96 | ### 3. Temporary Ban 97 | 98 | **Community Impact**: A serious violation of community standards, including 99 | sustained inappropriate behavior. 100 | 101 | **Consequence**: A temporary ban from any sort of interaction or public 102 | communication with the community for a specified period of time. No public or 103 | private interaction with the people involved, including unsolicited interaction 104 | with those enforcing the Code of Conduct, is allowed during this period. 105 | Violating these terms may lead to a permanent ban. 106 | 107 | ### 4. Permanent Ban 108 | 109 | **Community Impact**: Demonstrating a pattern of violation of community 110 | standards, including sustained inappropriate behavior, harassment of an 111 | individual, or aggression toward or disparagement of classes of individuals. 112 | 113 | **Consequence**: A permanent ban from any sort of public interaction within the 114 | community. 115 | 116 | ## Attribution 117 | 118 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], 119 | version 2.1, available at 120 | [https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1]. 121 | 122 | Community Impact Guidelines were inspired by 123 | [Mozilla's code of conduct enforcement ladder][Mozilla CoC]. 124 | 125 | For answers to common questions about this code of conduct, see the FAQ at 126 | [https://www.contributor-covenant.org/faq][FAQ]. Translations are available at 127 | [https://www.contributor-covenant.org/translations][translations]. 128 | 129 | [homepage]: https://www.contributor-covenant.org 130 | [v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html 131 | [Mozilla CoC]: https://github.com/mozilla/diversity 132 | [FAQ]: https://www.contributor-covenant.org/faq 133 | [translations]: https://www.contributor-covenant.org/translations 134 | -------------------------------------------------------------------------------- /doc/CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributor Guide 🤗 2 | 3 | We welcome various sorts of contributions to NNSmith, 4 | including reporting an [issue](https://github.com/ise-uiuc/nnsmith/issues) or submitting a [PR](https://github.com/ise-uiuc/nnsmith/pulls). 5 | 6 | ## Reporting an issue 7 | 8 | We appreciate developers to report current limitations of NNSmith on the [GitHub issue tracking system](https://github.com/ise-uiuc/nnsmith/issues), 9 | including but not limited to: 10 | 11 | 1. **Bug reports**: unexpected behaviors of NNSmith (e.g., a flaw of the operator rule/specification) 12 | 2. **Feature requests**: tell us what could be the promising feature to make NNSmith stronger! 13 | 14 | ## Submitting a PR 15 | 16 | The general flow for submitting a PR (Pull Request): 17 | 18 | > (Optional) Submit an issue to talk about the PR (necessary for introducing new features); 19 | 20 | 1. [Fork](https://github.com/ise-uiuc/nnsmith/fork) NNSmith; 21 | 2. Pull your fork: `git clone git@github.com:$[YOUR_FORK]/nnsmith.git`; 22 | 3. `cd nnsmith && export PYTHONPATH=$PYTHONPATH:$(pwd)` 23 | 4. Coding NNSmith! Then commit and push your code! 24 | 5. Submit a PR [here](https://github.com/ise-uiuc/nnsmith/pulls); 25 | 6. Code review; 26 | 7. Merge! 27 | 28 | ### Do I need to submit an issue before PR? 29 | 30 | - **No**: minor or nice-to-have cases such as typo fixes and bug fixes. 31 | - **Yes**: new features (e.g., extending new backend) and fundamental changes. 32 | 33 | ### Will my contributions be rejected? 34 | 35 | Oftentimes not, rare cases yes (that's why it is suggested to submit an issue for discussion first). 36 | 37 | **S-sized contributions** are oftentimes easy-to-accept, including bug/typo fixes, CI improvements, test-case improvements, etc. 38 | as long as it is beneficial and satisfies the properties in the "General coding guidance" section. 39 | 40 | **M-sized contributions** such as extending new front-ends/backends/fuzzing strategies/etc. are welcome as well 41 | -- as long as it shows an edge in improvements. 42 | However, for maintainability, it could be moved to the temporary "contrib" folder if it is non-trivial/unclear for being well-maintained. 43 | For example, let's say we supported backend "X" in the "contrib" folder and started to submitting bug reports to the "X" community. 44 | Later on, if "X" community is found to be not interested fixing bugs 45 | -- we don't have to support "X" as backend and consequently we can just drop it. 46 | 47 | **L-sized contributions** are those that conflicting the fundamental designs and goals of NNSmith. 48 | For example, NNSmith is fundamentally model generator, and it too much for it to support, for example, "distributed training". 49 | As a result, such changes might not be accepted unless there is a compelling justification 50 | -- but NNSmith is under Apache-2.0 -- you can always make it in the way you like via a fork :). 51 | Of course, some L-sized contributions can still possibly accepted, 52 | such as improving the operator specification or developing a more promising intermediate representation than GraphIR, 53 | as long as we agree on that the benefits (over the efforts) are unquestionable. 54 | 55 | ## General coding guidance 56 | 57 | ### `pre-commit` 58 | 59 | [`pre-commit`](https://pre-commit.com/) is a convenient tool to check and format your code while committing codes. 60 | 61 | To set-up pre-commit: 62 | 63 | ```shell 64 | pip install -r requirements/dev.txt 65 | pre-commit install 66 | ``` 67 | 68 | Now it will run checking and auto-formatting while you commit: 69 | 70 | ```shell 71 | git commit ... 72 | # if [NOTHING HAPPENS], you are good to go; 73 | # if [IT FAILS], the auto-formatting is automatically applied; 74 | # you just need to check, `git add` these changes and re-commit. 75 | ``` 76 | 77 | ### Testing 78 | 79 | If applicable (e.g., adding a new backend), add a few tests to validate your implementation. Examples can be found: 80 | 81 | 1. [Python unit-tests](https://github.com/ise-uiuc/nnsmith/tree/main/tests); 82 | 2. [End-to-end testing](https://github.com/ise-uiuc/nnsmith/blob/main/.github/workflows/ci.yaml); 83 | 84 | To run the Python tests: 85 | 86 | ```shell 87 | # env of torch & tf (and others) will conflict so split their unit tests. 88 | pytest tests/core -s 89 | pytest tests/torch -s 90 | pytest tests/tensorflow -s 91 | pytest tests/onnxruntime -s 92 | pytest tests/tvm -s 93 | pytest tests/tensorrt -s 94 | ``` 95 | 96 | ### Simple code 97 | 98 | > “Simplicity is the prerequisite for reliability.” - Edsger W. Dijkstra 99 | 100 | Maintaining code is hard, esp. when 101 | (i) initial code owners are not available; and 102 | (ii) the code is too complicated to be understood/modified. 103 | As a result, contributors are recommended to write simple code: 104 | (i) easy-to-understand; 105 | (ii) well-organized and easy-to-extend; 106 | (iii) well-documented if the concept is tricky; 107 | and (iv) avoiding changes that brings low improvement over high complexity. 108 | 109 | For example, the complexity of test-case structure is non-trivial in NNSmith; 110 | consequently, initial maintainers spent some amount of effort to make it systematically structured, 111 | so that it will be easier-to-use and extend. 112 | (@ganler: I know it could be boring, but it is indeed important for a long-live project.) 113 | 114 | ![](https://gist.github.com/ganler/bdf7e867e57c96e8c09ff31cb0b90a1f/raw/4667ad9b7dcb0b77cb722e7025402105560ebf41/datastructure.png) 115 | 116 | There are a few more concrete terms to consider: 117 | 118 | 1. Try not to introduce new dependencies: 119 | - If we only need "one" function from the prospective dependency, implement it on our own if possible; 120 | - If we have to use, try to consider "reliable" ones first. For example, those have been tested by millions of developers (such as NumPy). 121 | 2. Avoid bring data files in the repository -- it will bloat the codebase, making it harder to distribute. 122 | - If it is a picture, upstream that to gist or other "storage" repos and use an URL for it. 123 | - If it is some configuration file or data file, using script to re-generate it (if easy) or we distribute that on ["Releases"](https://github.com/ise-uiuc/nnsmith/releases) (if large). 124 | -------------------------------------------------------------------------------- /equality_saturation_helper/fuzz_loop.py: -------------------------------------------------------------------------------- 1 | # this script is used as a parent process for mitigating the memory leak of DL frameworks/compilers 2 | import logging 3 | import subprocess 4 | import os 5 | import argparse 6 | import time 7 | import shutil 8 | from subprocess import TimeoutExpired 9 | import signal 10 | import psutil 11 | 12 | TIME_FOR_SUBPROCESS = 1850 13 | 14 | 15 | # PARALLEL = 5 16 | 17 | 18 | def clean_per_circle(): 19 | inductor_cache_path = "/dev/shm/torchinductor" 20 | # shm_usage = psutil.disk_usage('/dev/shm') 21 | # if shm_usage.free > 10 * 1024 * 1024 * 1024: 22 | # os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache_path 23 | # else: 24 | # if "TORCHINDUCTOR_CACHE_DIR" in os.environ: 25 | # del os.environ["TORCHINDUCTOR_CACHE_DIR"] 26 | if os.path.exists("/tmp/torchinductor_root"): 27 | shutil.rmtree("/tmp/torchinductor_root") 28 | if os.path.exists("/root/.cache/hidet"): 29 | shutil.rmtree("/root/.cache/hidet") 30 | if os.path.exists(inductor_cache_path): 31 | shutil.rmtree(inductor_cache_path) 32 | 33 | 34 | def kill_a_proc(proc: subprocess.Popen) -> int: 35 | return_code = proc.poll() 36 | start_time = time.time() 37 | while return_code is None: # it is still alive 38 | if time.time() - start_time < 2: 39 | proc.terminate() 40 | else: 41 | proc.send_signal(signal.SIGKILL) 42 | return_code = proc.poll() 43 | print(f"spend {time.time() - start_time} second to kill the proc") 44 | return return_code 45 | 46 | 47 | def main(): 48 | parser = argparse.ArgumentParser(description="fuzzing loop helper") 49 | parser.add_argument("-o", "--output_path", type=str, help="directory path of output path", required=True) 50 | parser.add_argument("-s", "--script_path", type=str, help="path of fuzzing script", required=True) 51 | parser.add_argument("-m", "--model_type", type=str, help="torch/onnx/tensorflow", required=True) 52 | parser.add_argument("-b", "--backend_type", type=str, help="backend", required=True) 53 | parser.add_argument("-p", "--parallel", type=int, help="how many cores", required=True) 54 | parser.add_argument("-t", "--backend_target", type=str, help="cpu/gpu", default="cpu") 55 | parser.add_argument("-r", "--running-time", type=int, help="time for running (s), -1 for never stop, default -1", 56 | default=-1) 57 | parser.add_argument("-a", "--appended_args", type=str, help="appended arguments for the script", default="") 58 | args = parser.parse_args() 59 | script_path = args.script_path 60 | output_path = args.output_path 61 | model_type = args.model_type 62 | backend_type = args.backend_type 63 | parallel = args.parallel 64 | backend_target = args.backend_target 65 | running_time = args.running_time 66 | appended_args = [i for i in args.appended_args.split(" ")] if args.appended_args != "" else [] 67 | if os.path.exists(output_path): 68 | logging.error(f"{output_path} already exists, please remove it before we can start the fuzzing process") 69 | return 70 | 71 | os.makedirs(output_path, exist_ok=True) 72 | 73 | time_for_start_all = time.time() 74 | session_id = 0 75 | while True: 76 | clean_per_circle() 77 | if running_time != -1 and time.time() - time_for_start_all > running_time: 78 | print("fuzzing is over due to timeout") 79 | break 80 | 81 | print(f"new cycle, session id start with {session_id}, parallel num is {parallel}") 82 | procs = [] 83 | start_time = time.time() 84 | for _ in range(parallel): 85 | fuzzer_output_dir = os.path.join(output_path, f"fuzz_report_{session_id}") 86 | log_output_path = os.path.join(output_path, f"fuzz_log_{session_id}") 87 | with open(log_output_path, "w+") as f: 88 | cmd = ["python3", script_path, f"fuzz.time={TIME_FOR_SUBPROCESS - 50}s", "cmp.oracle=null", 89 | "cmp.oracle=null", 90 | f"model.type={model_type}", 91 | f"backend.type={backend_type}", f"backend.target={backend_target}", 92 | f"fuzz.root={fuzzer_output_dir}", 93 | "debug.viz=true", "mgen.max_nodes=10"] + appended_args 94 | session_id += 1 95 | proc = subprocess.Popen(cmd, stdout=f, stderr=f) 96 | procs.append(proc) 97 | 98 | print(f"wait for the termination of {len(procs)} processes, time: {time.asctime(time.localtime(time.time()))}") 99 | used = [False for _ in range(len(procs))] 100 | jump_out = False 101 | while True: 102 | if used.count(True) == len(used): 103 | break 104 | 105 | if (time.time() - start_time) > TIME_FOR_SUBPROCESS: 106 | print(f"time to terminate. time: {time.asctime(time.localtime(time.time()))}") 107 | 108 | for index, proc in enumerate(procs): 109 | if used[index]: 110 | continue 111 | 112 | if jump_out: 113 | return_code = kill_a_proc(proc) 114 | print(f"process {index} has been terminated due to jump out; exit code: {return_code}") 115 | used[index] = True 116 | 117 | try: 118 | code = proc.wait(5) 119 | print(f"process {index} has been normally terminated; exit code: {code}") 120 | used[index] = True 121 | except TimeoutExpired as _: 122 | if (time.time() - start_time) > TIME_FOR_SUBPROCESS and proc.poll() is None: 123 | code = kill_a_proc(proc) 124 | print(f"process {index} has been force terminated; exit code: {code}") 125 | used[index] = True 126 | except BaseException as e: 127 | print(f"terminated by {e}, start killing the processes") 128 | jump_out = True 129 | 130 | print("all processes are terminated") 131 | if jump_out: 132 | break 133 | print("all done") 134 | 135 | 136 | if __name__ == '__main__': 137 | main() 138 | -------------------------------------------------------------------------------- /experiments/legacy/lemon_tf2onnx.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pathlib import Path 3 | import argparse 4 | from multiprocessing import Process 5 | import time 6 | 7 | import tensorflow as tf 8 | from tensorflow import keras 9 | from tensorflow.python.framework.convert_to_constants import ( 10 | convert_variables_to_constants_v2, 11 | ) 12 | import tf2onnx 13 | from tqdm import tqdm 14 | import numpy as np 15 | 16 | from nnsmith.util import mkdir 17 | 18 | tf.get_logger().setLevel("WARNING") # Tensorflow made quiet. 19 | 20 | 21 | def analyze_inputs_outputs(graph): 22 | ops = graph.get_operations() 23 | outputs_set = set(ops) 24 | inputs = [] 25 | for op in ops: 26 | if len(op.inputs) == 0 and op.type == "Placeholder": 27 | inputs.append(op) 28 | else: 29 | for input_tensor in op.inputs: 30 | if input_tensor.op in outputs_set: 31 | outputs_set.remove(input_tensor.op) 32 | outputs = list(outputs_set) 33 | # Control nodes shall not be considered. 34 | # input like: "import/x" -> x 35 | # output like: "import/Identity", "import/Identity_1" -> Identity, Identity_1 36 | inputs = [x.name.split("/")[-1] for x in inputs if "_control_node" not in x.name] 37 | outputs = [x.name.split("/")[-1] for x in outputs if "_control_node" not in x.name] 38 | return (inputs, outputs) 39 | 40 | 41 | def keras2tf(model): 42 | full_model = tf.function(lambda x: model(x)) 43 | freeze_shape = model.inputs[0].shape 44 | 45 | shape_list = [] 46 | for v in freeze_shape: 47 | try: 48 | shape_list.append(int(v)) 49 | except TypeError as e: 50 | shape_list.append(1) 51 | 52 | full_model = full_model.get_concrete_function( 53 | tf.TensorSpec(tf.TensorShape(shape_list), model.inputs[0].dtype) 54 | ) 55 | frozen_func = convert_variables_to_constants_v2(full_model) 56 | 57 | return frozen_func.graph.as_graph_def() 58 | 59 | 60 | def convert_tf2onnx(from_path, to_path): 61 | def custom_objects(): 62 | def no_activation(x): 63 | return x 64 | 65 | def leakyrelu(x): 66 | return keras.activations.relu(x, alpha=0.01) 67 | 68 | objects = {} 69 | objects["no_activation"] = no_activation 70 | objects["leakyrelu"] = leakyrelu 71 | return objects 72 | 73 | model = keras.models.load_model(from_path, custom_objects=custom_objects()) 74 | shape_list = [] 75 | for v in model.inputs[0].shape: 76 | try: 77 | shape_list.append(int(v)) 78 | except TypeError as e: 79 | shape_list.append(1) 80 | 81 | graph_def = keras2tf(model) 82 | 83 | with tf.Graph().as_default() as graph: 84 | tf.import_graph_def(graph_def) 85 | inps, outs = analyze_inputs_outputs(graph) 86 | 87 | tf2onnx.convert.from_graph_def( 88 | graph_def, 89 | input_names=[inp + f":{i}" for i, inp in enumerate(inps)], 90 | output_names=[out + f":{i}" for i, out in enumerate(outs)], 91 | opset=13, 92 | output_path=to_path, 93 | ) 94 | 95 | 96 | if __name__ == "__main__": 97 | parser = argparse.ArgumentParser() 98 | parser.add_argument( 99 | "--lemon_output_dir", type=str, required=True, help="Path to the folder." 100 | ) 101 | parser.add_argument("--onnx_dir", type=str, required=True) 102 | parser.add_argument("--tlimit", type=int, default=4 * 60 * 60) 103 | args = parser.parse_args() 104 | 105 | mkdir(args.onnx_dir) 106 | 107 | # FORMAT: {generation time cost in seconds}, {model relative path} 108 | # MUST RANK by GENERATION ORDER. 109 | config_file = open(os.path.join(args.onnx_dir, "gentime.csv"), "w") 110 | 111 | time_list = [] 112 | file_list = [] 113 | for file in Path(args.lemon_output_dir).rglob("*/*.h5"): 114 | file_list.append(file) 115 | time_list.append(file.stat().st_mtime) 116 | 117 | assert len(file_list) > 0, "No files found in the folder!" 118 | print(f"{len(file_list)} files found in the folder.") 119 | 120 | time_arr = np.array(time_list) 121 | time_arr -= time_arr.min() 122 | idx = time_arr.argsort() 123 | 124 | time_span = time_arr.max() - time_arr.min() 125 | print( 126 | f"|T_last - T_first| = {time_span:.2f}s = {time_span / 60:.2f}min = {time_span / 3600:.2f}h" 127 | ) 128 | assert ( 129 | time_span / 3600 < 25 130 | ), "Are you sure your steps are correct? The time span is too long..." 131 | 132 | time_diffs = np.diff(np.sort(time_arr)) 133 | 134 | cvt_start_time = time.time() 135 | 136 | for i in tqdm(range(len(file_list))): 137 | ranked_idx = idx[i] 138 | from_path = file_list[ranked_idx] 139 | to_name = os.path.split(from_path)[-1] + ".onnx" 140 | to_path = os.path.join(args.onnx_dir, to_name) 141 | try: 142 | tstart = time.time() 143 | # might crash... 144 | p = Process(target=convert_tf2onnx, args=(from_path, to_path)) 145 | p.start() 146 | p.join() 147 | 148 | assert p.exitcode is not None 149 | if p.exitcode != 0: 150 | raise RuntimeError(f"Conversion failed with exit code {p.exitcode}!") 151 | 152 | # We evaluate the end2end efficiency so we count everything. 153 | cvt_time = time.time() - tstart 154 | if i == 0: 155 | config_file.write(f"{time_diffs.mean() + cvt_time},{to_name}\n") 156 | else: 157 | config_file.write(f"{time_diffs[i - 1] + cvt_time},{to_name}\n") 158 | except Exception as e: 159 | # Models generated by LEMON are sometimes not valid. 160 | # We evaluate the end2end efficiency so we count everything. 161 | cvt_time = time.time() - tstart 162 | if i == 0: 163 | config_file.write(f"{time_diffs.mean() + cvt_time},FAILURE\n") 164 | else: 165 | config_file.write(f"{time_diffs[i - 1] + cvt_time},FAILURE\n") 166 | print(e) 167 | config_file.flush() 168 | if time_diffs[:i].sum() + (time.time() - cvt_start_time) > args.tlimit: 169 | print(f"Time limit reached. {i+1} models converted.") 170 | break 171 | config_file.close() 172 | -------------------------------------------------------------------------------- /experiments/legacy/batch_eval.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import random 4 | import numpy as np 5 | import pickle 6 | 7 | from nnsmith.error import IncorrectResult 8 | from nnsmith.backends import BackendFactory, mk_factory 9 | from nnsmith.difftest import assert_allclose 10 | from nnsmith.util import gen_one_input 11 | from nnsmith.fuzz import simple_bug_report 12 | 13 | 14 | def mcov_write(path): 15 | if path: 16 | with open(path + ".pkl", "wb") as f: 17 | pickle.dump(omax_fac._coverage_install().get_hitmap(), f) 18 | 19 | 20 | def verify(backend_name, oracle_name, predicted, oracle=None): 21 | try: 22 | if oracle is not None: 23 | assert_allclose(predicted, oracle, backend_name, oracle_name) 24 | except IncorrectResult as e: 25 | return e 26 | return None 27 | 28 | 29 | if __name__ == "__main__": 30 | import argparse 31 | 32 | parser = argparse.ArgumentParser() 33 | parser.add_argument( 34 | "--models", nargs="+", required=True, help="List to ONNX model paths" 35 | ) 36 | parser.add_argument( 37 | "--backend", type=str, help="One of ort, trt, tvm, and xla", required=True 38 | ) 39 | parser.add_argument("--device", type=str, default="cpu") 40 | parser.add_argument( 41 | "--memcov", type=str, default=None, help="Path to store memcov." 42 | ) 43 | parser.add_argument( 44 | "--seed", type=int, default=233, help="to generate random input data" 45 | ) 46 | parser.add_argument("--fuzz_max_nodes", type=int, help="parameter from fuzzer") 47 | parser.add_argument("--fuzz_seed", type=int, help="seed parameter from fuzzer") 48 | parser.add_argument("--fuzz_report_folder", type=str, help="parameter from fuzzer") 49 | parser.add_argument( 50 | "--clean_after_eval", action="store_true", help="rm models/oracle after eval" 51 | ) 52 | # add fuzz_timeout? 53 | args = parser.parse_args() 54 | 55 | if args.fuzz_report_folder is None: 56 | print( 57 | "[WARNING] Bug report is not enabled as fuzzer parameters are not provided.", 58 | file=sys.stderr, 59 | ) 60 | 61 | # Set global seed 62 | random.seed(args.seed) 63 | np.random.seed(args.seed) 64 | 65 | omax_fac = mk_factory(args.backend, args.device, optmax=True) 66 | omin_fac = mk_factory(args.backend, args.device, optmax=False) 67 | 68 | if args.memcov: 69 | assert omax_fac._coverage_install().get_now() is not None, "Memcov unavailable!" 70 | 71 | for i, path in enumerate(args.models): 72 | print(f"-> {path}", flush=True, file=sys.stderr) 73 | onnx_model = BackendFactory.get_onnx_proto(path) 74 | oracle_path = path.replace(".onnx", ".pkl") 75 | 76 | num_must_valid = False 77 | if os.path.exists(oracle_path): 78 | with open(oracle_path, "rb") as f: 79 | res = pickle.load(f) 80 | eval_inputs = res[0] 81 | torch_outputs = res[1] 82 | if len(res) == 3: 83 | num_must_valid = res[2] 84 | else: 85 | print(f"No oracle found for model `{path}`", file=sys.stderr) 86 | input_spec, onames = BackendFactory.analyze_onnx_io(onnx_model) 87 | eval_inputs = gen_one_input(input_spec, 1, 2) 88 | torch_outputs = None # No oracle. 89 | 90 | to_repro = f"python nnsmith/graph_gen.py --max_nodes {args.fuzz_max_nodes} --seed {args.fuzz_seed} --viz_graph" 91 | 92 | # Convenience function to report bugs. 93 | def report_bug(category, exception): 94 | simple_bug_report( 95 | report_folder=args.fuzz_report_folder, 96 | buggy_onnx_path=path, 97 | oracle_path=oracle_path, 98 | message=to_repro + "\n" + str(exception), 99 | bug_type=category + type(exception).__name__, 100 | backend=args.backend, 101 | ) 102 | 103 | try: 104 | omax_exec = omax_fac.make_backend(onnx_model) 105 | pred_omax = omax_exec(eval_inputs) 106 | # We compare results iff NaN/Inf-free as we only consider confident bugs. 107 | if num_must_valid and all( 108 | np.isfinite(v).all() for v in torch_outputs.values() 109 | ): 110 | # FIXME: We assume the oracle is generated by PyTorch (as `fuzz.py`). 111 | err_tch_omax = verify( 112 | args.backend + "-omax", "PyTorch", pred_omax, torch_outputs 113 | ) 114 | if err_tch_omax: 115 | if omin_fac is None: 116 | # For TensorRT, there's no O0 option. 117 | report_bug("torch-omax-", err_tch_omax) 118 | else: 119 | # RESULT INCONSISTENCY. But need to see if we can narrow it down to an optimzation bug. 120 | omin_exec = omin_fac.make_backend(onnx_model) 121 | err_omax_omin = verify( 122 | args.backend + "-omax", 123 | args.backend + "-omin", 124 | pred_omax, 125 | omin_exec(eval_inputs), 126 | ) 127 | if err_omax_omin: 128 | # We can confirm this is an optimization bug as O[opt-max] != O[opt-min] 129 | report_bug("omin-omax-", err_omax_omin) 130 | else: 131 | # O[opt-max] == O[opt-min]: even opt-min shows the inconsistency. 132 | # So the bug might be a fault of `PyTorch` or other parts in the compiler... Human assistant is needed. 133 | report_bug("torch-omax-", err_tch_omax) 134 | except Exception as e: 135 | # Well, some error occured (compiler internal err | nnsmith impl error). 136 | report_bug("cie-", e) # Compiler internal error. 137 | 138 | # remove after the model is tested. useful for locating the crashed model in the batch. 139 | if args.clean_after_eval: 140 | if os.path.exists(path): 141 | os.unlink(path) 142 | if os.path.exists(oracle_path): 143 | os.unlink(oracle_path) 144 | 145 | mcov_write(args.memcov) 146 | 147 | mcov_write(args.memcov) 148 | -------------------------------------------------------------------------------- /experiments/legacy/plot_inp_search_merge.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import argparse 4 | 5 | import matplotlib.pyplot as plt 6 | import matplotlib.patches as mpatches 7 | import pandas as pd 8 | import numpy as np 9 | 10 | 11 | SMALL_SIZE = 10 12 | MEDIUM_SIZE = 13 13 | BIGGER_SIZE = 18 14 | 15 | plt.rc("font", size=SMALL_SIZE) # controls default text sizes 16 | plt.rc("axes", titlesize=MEDIUM_SIZE) # fontsize of the axes title 17 | plt.rc("axes", labelsize=MEDIUM_SIZE) # fontsize of the x and y labels 18 | plt.rc("xtick", labelsize=SMALL_SIZE) # fontsize of the tick labels 19 | plt.rc("ytick", labelsize=SMALL_SIZE) # fontsize of the tick labels 20 | plt.rc("legend", fontsize=MEDIUM_SIZE) # legend fontsize 21 | plt.rc("figure", titlesize=BIGGER_SIZE) # fontsize of the figure title 22 | 23 | plt.rcParams.update({"text.usetex": True}) 24 | 25 | if __name__ == "__main__": 26 | parser = argparse.ArgumentParser() 27 | parser.add_argument("--root", type=str, nargs="+", required=True) 28 | parser.add_argument("--output", type=str, default="results") 29 | parser.add_argument( 30 | "--rm_slowest", action="store_true", help="Remove the slowest run" 31 | ) 32 | args = parser.parse_args() 33 | 34 | REGEX_PATTERN = "(\d+)-model-(\d+)-node-exp" 35 | node_sizes = [] 36 | 37 | # Plot data 38 | # X: Time 39 | # Y: Succ. Rate 40 | 41 | def sort_by_time(time, succ_rate): 42 | time = np.array(time) 43 | succ_rate = np.array(succ_rate) 44 | sort_idx = time.argsort() 45 | return time[sort_idx], succ_rate[sort_idx] 46 | 47 | sampling_res = [] 48 | grad_res = [] 49 | proxy_res = [] 50 | 51 | for nsize_folder in args.root: 52 | res = re.match(REGEX_PATTERN, nsize_folder) 53 | n_model, n_nodes = res.groups() 54 | node_sizes.append(int(n_nodes)) 55 | 56 | sampling_time = [] 57 | sampling_succ_rate = [] 58 | 59 | grad_time = [] 60 | grad_succ_rate = [] 61 | 62 | proxy_time = [] 63 | proxy_succ_rate = [] 64 | 65 | for f in os.listdir(nsize_folder): 66 | if f.endswith(".csv") and f != "model_info.csv": 67 | data = pd.read_csv(os.path.join(nsize_folder, f)) 68 | 69 | # Do not count the slowest iteration (1st iter usually) 70 | # as initialization takes some time. 71 | last_idx = -2 if args.rm_slowest else -1 72 | 73 | idx = data["sampling-time"].to_numpy().argsort()[:last_idx] 74 | sampling_time.append(data["sampling-time"][idx].mean()) 75 | sampling_succ_rate.append(data["sampling-succ"][idx].mean()) 76 | 77 | idx = data["grad-time"].to_numpy().argsort()[:last_idx] 78 | grad_time.append(data["grad-time"][idx].mean()) 79 | grad_succ_rate.append(data["grad-succ"][idx].mean()) 80 | 81 | idx = data["proxy-time"].to_numpy().argsort()[:last_idx] 82 | proxy_time.append(data["proxy-time"][idx].mean()) 83 | proxy_succ_rate.append(data["proxy-succ"][idx].mean()) 84 | elif f == "model_info.csv": 85 | data = pd.read_csv(os.path.join(nsize_folder, f)) 86 | gentime = data["gen_time"] 87 | print(f"{gentime.mean()=}, {gentime.min()=}, {gentime.max()=}") 88 | 89 | # sort succ rate by time 90 | sampling_res.append(sort_by_time(sampling_time, sampling_succ_rate)) 91 | grad_res.append(sort_by_time(grad_time, grad_succ_rate)) 92 | proxy_res.append(sort_by_time(proxy_time, proxy_succ_rate)) 93 | 94 | fig, ax = plt.subplots(figsize=(8, 3.2), constrained_layout=True) 95 | 96 | colors = ["dodgerblue", "violet", "green"] # ['b', 'r', 'g'] 97 | markers = ["1", ".", "+"] 98 | markercolor = "k" 99 | markersize = 10 100 | markeredgewidth = 1.2 101 | lw = 1.5 102 | 103 | max_time = 0 104 | 105 | for i in range(3): 106 | c = colors[i] 107 | alpha = 1 - 0.36 * i 108 | 109 | sampling_time, sampling_succ_rate = sampling_res[i] 110 | grad_time, grad_succ_rate = grad_res[i] 111 | proxy_time, proxy_succ_rate = proxy_res[i] 112 | 113 | # self *= 1000: sec -> milli 114 | sampling_time *= 1000 115 | grad_time *= 1000 116 | proxy_time *= 1000 117 | 118 | ax.plot( 119 | proxy_time, 120 | proxy_succ_rate, 121 | marker=markers[0], 122 | markeredgecolor=markercolor, 123 | markersize=markersize, 124 | markeredgewidth=markeredgewidth, 125 | linestyle="--", 126 | color=c, 127 | lw=lw, 128 | ) 129 | ax.plot( 130 | grad_time, 131 | grad_succ_rate, 132 | marker=markers[1], 133 | markeredgecolor=markercolor, 134 | markersize=markersize, 135 | markeredgewidth=markeredgewidth, 136 | linestyle="--", 137 | color=c, 138 | lw=lw, 139 | ) 140 | ax.plot( 141 | sampling_time, 142 | sampling_succ_rate, 143 | marker=markers[2], 144 | markeredgecolor=markercolor, 145 | markersize=markersize, 146 | markeredgewidth=markeredgewidth, 147 | linestyle="--", 148 | color=c, 149 | lw=lw, 150 | ) 151 | 152 | max_time = max(max_time, sampling_time.max()) 153 | 154 | ax.grid(True, linestyle=":", linewidth=0.5, alpha=0.5) 155 | 156 | lines = ax.get_lines() 157 | legend1 = plt.legend( 158 | [lines[i] for i in [0, 1, 2]], 159 | ["Gradient (Proxy Deriv.)", "Gradient", "Sampling"], 160 | loc="upper right", 161 | title="Searching Method", 162 | ) 163 | 164 | patches = [] 165 | for i in range(3): 166 | patches.append(mpatches.Patch(color=colors[i], label=node_sizes[i])) 167 | legend2 = plt.legend( 168 | handles=patches, 169 | loc="center right", 170 | bbox_to_anchor=(1, 0.36), 171 | title="Model Size", 172 | ) 173 | 174 | ax.add_artist(legend1) 175 | ax.add_artist(legend2) 176 | 177 | ax.set_yticks(np.arange(0.6, 1.1, 0.1)) 178 | ax.set_ylim(0.6, 1.0) 179 | 180 | ax.set_xticks(np.arange(0, 31, 5)) 181 | ax.set_xlim(0, max_time + 0.5) 182 | 183 | ax.set_xlabel("Avg. Searching Time (millisecond)", fontweight="bold") 184 | ax.set_ylabel("Success Rate", fontweight="bold") 185 | 186 | plt.savefig(os.path.join(args.output, f"input-search-merge.pdf")) 187 | plt.savefig(os.path.join(args.output, f"input-search-merge.png")) 188 | -------------------------------------------------------------------------------- /nnsmith/backends/tflite.py: -------------------------------------------------------------------------------- 1 | from os import PathLike 2 | from typing import Dict, List 3 | 4 | import numpy as np 5 | import tensorflow as tf # type: ignore 6 | from multipledispatch import dispatch 7 | 8 | from nnsmith.abstract.dtype import DType 9 | from nnsmith.backends.factory import BackendCallable, BackendFactory 10 | from nnsmith.materialize.tensorflow import TFModel, TFNetCallable 11 | 12 | 13 | class TFLiteRunner: 14 | def __init__(self, tfnet_callable: TFNetCallable) -> None: 15 | self.tfnet_callable = tfnet_callable 16 | 17 | def __call__(self, input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: 18 | return {k: np.array(v) for k, v in self.tfnet_callable(**input).items()} 19 | 20 | 21 | class TFLite(BackendFactory): 22 | """Factory to build TFLite backend. 23 | Convertion Graph: 24 | TFModel & one concrete function 25 | v TFLiteConverter with configs 26 | TFLite model content in bytes -- (f.write) --> model.tflite file 27 | v tf.lite.Interpreter <-- (f.read) <--- 28 | TFLite Python Callable 29 | """ 30 | 31 | def __init__(self, target, optmax, **kwargs) -> None: 32 | # https://github.com/tensorflow/tensorflow/issues/34536#issuecomment-565632906 33 | # TFLite doesn't support NVIDIA GPU. 34 | assert target != "cuda" 35 | super().__init__(target, optmax, **kwargs) 36 | 37 | @property 38 | def system_name(self) -> str: 39 | return "tflite" 40 | 41 | @property 42 | def version(self) -> str: 43 | return tf.__version__ 44 | 45 | @dispatch(TFModel) 46 | def make_backend( 47 | self, 48 | model: TFModel, 49 | ) -> BackendCallable: 50 | """Create TFLite callable from a concrete function. 51 | TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`. 52 | Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions 53 | """ 54 | return self.make_backend_from_content(self.make_content(model)) 55 | 56 | def make_backend_from_path( 57 | self, 58 | path: PathLike, 59 | ) -> BackendCallable: 60 | """Create TFLite callable from path of a TF SavedModel. 61 | Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_saved_model 62 | """ 63 | return self.make_backend_from_content( 64 | self.make_content(path), 65 | ) 66 | 67 | @dispatch(bytes) 68 | def make_backend_from_content(self, content: bytes) -> BackendCallable: 69 | # Ref: https://www.tensorflow.org/api_docs/python/tf/lite/Interpreter 70 | interpreter = tf.lite.Interpreter(model_content=content) 71 | return TFLiteRunner(interpreter.get_signature_runner()) 72 | 73 | @dispatch(TFModel) 74 | def make_content( 75 | self, 76 | model: TFModel, 77 | ) -> bytes: 78 | """Create TFLite content from a concrete function. 79 | TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`. 80 | Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions 81 | """ 82 | converter = tf.lite.TFLiteConverter.from_concrete_functions( 83 | funcs=[model.concrete_net()], 84 | trackable_obj=model.net, 85 | ) 86 | return self._tflite_content_from_converter(converter) 87 | 88 | def make_content_with_func( 89 | self, 90 | model: TFModel, 91 | concrete_func: TFNetCallable, 92 | ) -> bytes: 93 | """Create TFLite content from a concrete function. 94 | TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`. 95 | Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions 96 | """ 97 | converter = tf.lite.TFLiteConverter.from_concrete_functions( 98 | funcs=[concrete_func], 99 | trackable_obj=model.net, 100 | ) 101 | return self._tflite_content_from_converter(converter) 102 | 103 | def make_backend_with_func( 104 | self, 105 | model: TFModel, 106 | concrete_func: TFNetCallable, 107 | ) -> BackendCallable: 108 | """Create TFLite callable from a concrete function. 109 | TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`. 110 | Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions 111 | """ 112 | return self.make_backend_from_content( 113 | self.make_content_with_func(model, concrete_func) 114 | ) 115 | 116 | @dispatch(str) 117 | def make_content( 118 | self, 119 | path: PathLike, 120 | ) -> bytes: 121 | """Create TFLite content from path of a saved model. 122 | Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_saved_model 123 | """ 124 | converter = tf.lite.TFLiteConverter.from_saved_model(path) 125 | return self._tflite_content_from_converter(converter) 126 | 127 | def _tflite_content_from_converter( 128 | self, 129 | converter: tf.lite.TFLiteConverter, 130 | ) -> bytes: 131 | """Configure TFLite converter and create callable from it. 132 | converter configuarations: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#attributes_1 133 | """ 134 | # Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TargetSpec 135 | converter.target_spec.supported_ops = [ 136 | tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops. 137 | tf.lite.OpsSet.SELECT_TF_OPS, # enable TensorFlow ops. 138 | ] 139 | # Ref: https://www.tensorflow.org/api_docs/python/tf/lite/Optimize 140 | converter.optimizations = { 141 | tf.lite.Optimize.DEFAULT, 142 | tf.lite.Optimize.EXPERIMENTAL_SPARSITY, 143 | } 144 | # converter.allow_custom_ops = True 145 | tflite_bytes = converter.convert() 146 | return tflite_bytes 147 | 148 | def dump_backend(self, path: PathLike, content: bytes) -> None: 149 | with open(path, "wb") as f: 150 | f.write(content) 151 | 152 | def load_content(self, path: PathLike) -> bytes: 153 | with open(path, "rb") as f: 154 | return f.read() 155 | 156 | def load_backend(self, path: PathLike) -> BackendCallable: 157 | return self.make_backend_from_content(self.load_content(path)) 158 | 159 | @property 160 | def import_libs(self) -> List[str]: 161 | return ["import tensorflow as tf"] 162 | 163 | @classmethod 164 | def skip_dtypes(cls) -> List[DType]: 165 | # PyTorch-ONNX current does not support fully complex (e.g., scalar). 166 | # uint8 crash over Equal (revert after their fix). 167 | return [DType.uint8] 168 | --------------------------------------------------------------------------------