Add some TF ops and add XLA backend (ise-uiuc#40)

* feat: impl some forward functions for tf

* feat: add more dtypes

* fix: remove f16 from dtype_floats as it is not supported by many torch operators

* feat: add LocalRespNorm for TF

* feat: force CPU when making oracle with TF

* feat: add XLA as a backend

* feat: add tensorflow dump load test

* feat: try to add NHWCConv2d for TF

* test: add XLA test

* fix: disable autocast for keras layers

* fix: enable XLA compilation

* fix: typo

* chore: del commented code

* fix: disable some dtypes

* fix: disable some dtypes

* fix:reshape

* ci: complete tflite and xla test

* fix: tflite system name

* tweak tf logging and fix topset

* clean headers

* fix dtype reload

* hot fix: ci

* feat: add pattern to bug report folder

* fix: do not allow conv2d with valid padding to output tensor with dim=0

* refact: optimize logging for model gen

* pre-commit

* instruction for running pre-commit for all files; fix: error switch

* many fixes: supress absl, fix tf.conv spec, more logging levels

* pre-commit

* feat: dtype test with model gen only; align with testing

* feat: scm for auto versioning

* refact: use scm ver

* add coverage info

* show pytest log

* feat: optimize dependency

* fix: warn and relax pygraphviz

* skip gpu tests

* refact: disable pytest capture

* revert capture

Co-authored-by: ganler <[email protected]>

Author: Co1lin

.gitignore

@@ -13,6 +13,7 @@ tmp/
 *.profraw
 fuzz_report/
 nnsmith_output/
+nnsmith/_version.py
 
 # hydra
 outputs/

README.md

@@ -2,6 +2,18 @@
 
 [![](https://github.com/ise-uiuc/nnsmith/actions/workflows/ci.yaml/badge.svg)](https://github.com/ise-uiuc/nnsmith/actions/workflows/ci.yaml) [![](https://img.shields.io/pypi/v/nnsmith?color=g)](https://pypi.org/project/nnsmith/) [![](https://img.shields.io/pypi/l/nnsmith)](https://github.com/ise-uiuc/nnsmith/blob/main/LICENSE)
 
+## Backend-Model Support
+
+| Backend\Model | ONNX/PyTorch | TensorFlow |
+| ------------- | ------------ | ---------- |
+| TVM           | ✅            |            |
+| ONNXRuntime   | ✅            |            |
+| TensorRT      | ✅            |            |
+| TFLite        |              | ⚠️          |
+| XLA           |              | ⚠️          |
+
+✅: Supported; ⚠️: Beta support; Others are not supported yet -- Contributions are welcome!
+
 ## Quick Start
 
 <details><summary><b>Setting up graphviz for debugging</b> <i>[click to expand]</i></summary>
@@ -47,6 +59,7 @@ You can use `pre-commit` to simpify development:
 
 - `pip install -r requirements/dev.txt`;
 - `pre-commit install`;
+- `pre-commit` will run upon a commit; To explicitly run `pre-commit` for all files: `pre-commit run --all-files`.
 
 <details><summary><b>More notes</b> <i>[click to expand]</i></summary>
 <div>

doc/cli.md

@@ -66,3 +66,7 @@ nnsmith.dtype_test model.type="onnx" backend.type="onnxruntime"
 ```shell
 nnsmith.fuzz fuzz.time=30s fuzz.root=fuzz_report debug.viz=true
 ```
+
+## Misc
+
+TensorFlow logging can be very noisy. Use `TF_CPP_MIN_LOG_LEVEL=3` as environmental variable to depress that.

doc/log-and-err.md

@@ -6,17 +6,19 @@ We support the following logging "keys":
 
 - `fuzz`: fuzzing loop;
 - `mgen`: model generation;
+- `smt`: constraints in smt solving;
 - `exec`: model execution;
 - `viz`: graphviz visualization;
 - `dtest`: dtype_test;
+- `core`: seed setting, etc;
 
 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/)).
 
 ```shell
 # Show debug information related to `fuzz`:
 ${NNSMITH_CMD} hydra.verbose=fuzz
 # Show debug info for `fuzz` and `exec`:
-${NNSMITH_CMD} hydra.verbose=[fuzz,exec]
+${NNSMITH_CMD} hydra.verbose="[fuzz,exec]"
 ```
 
 #### Logging things into file

nnsmith/__init__.py

@@ -0,0 +1,4 @@
+try:
+    from nnsmith._version import __version__, __version_tuple__
+except ImportError:
+    pass

nnsmith/abstract/arith.py

@@ -181,3 +181,28 @@ def nnsmith_min(left, right):
         return min(left, right)
     left, right = align_bvs(left, right)
     return z3.If(nnsmith_le(left, right), left, right)
+
+
+def nnsmith_max(left, right):
+    if isinstance(left, int) and isinstance(right, int):
+        return min(left, right)
+    left, right = align_bvs(left, right)
+    return z3.If(nnsmith_ge(left, right), left, right)
+
+
+def nnsmith_and(left, right):
+    if isinstance(left, bool) and isinstance(right, bool):
+        return left and right
+    return z3.And(left, right)
+
+
+def nnsmith_or(left, right):
+    if isinstance(left, bool) and isinstance(right, bool):
+        return left or right
+    return z3.Or(left, right)
+
+
+def nnsmith_not(expr):
+    if isinstance(expr, bool):
+        return not expr
+    return z3.Not(expr)

nnsmith/abstract/dtype.py

@@ -116,25 +116,43 @@ def tensorflow(self) -> "tf.Dtype":
         import tensorflow as tf
 
         return {
+            DType.float16: tf.float16,
             DType.float32: tf.float32,
             DType.float64: tf.float64,
+            DType.int8: tf.int8,
+            DType.int16: tf.int16,
             DType.int32: tf.int32,
             DType.int64: tf.int64,
+            DType.complex64: tf.complex64,
+            DType.complex128: tf.complex128,
+            DType.bool: tf.bool,
         }[self]
 
     @staticmethod
     def from_tensorflow(dtype) -> "DType":
         import tensorflow as tf
 
         return {
+            tf.float16: DType.float16,
             tf.float32: DType.float32,
             tf.float64: DType.float64,
+            tf.int8: DType.int8,
+            tf.int16: DType.int16,
             tf.int32: DType.int32,
             tf.int64: DType.int64,
+            tf.complex64: DType.complex64,
+            tf.complex128: DType.complex128,
+            tf.bool: DType.bool,
         }[dtype]
 
 
-DTYPE_ALL = [DType.float32, DType.float64, DType.int32, DType.int64, DType.bool]
+DTYPE_ALL = [
+    DType.float32,
+    DType.float64,
+    DType.int32,
+    DType.int64,
+    DType.bool,
+]
 DTYPE_NON_BOOLS = [dtype for dtype in DTYPE_ALL if dtype != DType.bool]
 DTYPE_FLOATS = [DType.float32, DType.float64]
 DTYPE_INTS = [DType.int32, DType.int64]

nnsmith/abstract/op.py

@@ -1512,7 +1512,7 @@ def __init__(self, *target_shape):
         super().__init__()
         self.inp_ranks = [int_range(1, 4)]
         self.out_ranks = [(len(target_shape),)]
-        self.target_shape: List[Union[int, z3.ExprRef]] = target_shape
+        self.target_shape: List[Union[int, z3.ExprRef]] = list(target_shape)
 
     def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]:
         __MAX_SOLVE_SYMBOL__ = 8
@@ -1602,16 +1602,6 @@ def deduct_inp_ranks_and_dtype(
         return [(-1, out_abs_tensor[0].dtype)]
 
 
-@mark_materialize("core")
-class Flatten(Reshape):
-    num_var_param = None
-    # Inputs are target shape.
-
-    def __init__(self, dim0: Union[int, z3.ExprRef]):
-        super().__init__(dim0)
-        self.dim0 = dim0
-
-
 @mark_materialize("core")
 class Transpose(UnaryOpBase):
     in_dtypes = [(i,) for i in DTYPE_ALL]
@@ -1665,7 +1655,7 @@ class InterpBase(UnaryOpBase):
 
     def __init__(self, *size):
         super().__init__()
-        self.size = size
+        self.size = list(size)
         self.inp_ranks = [(len(size) + 2,)]
         self.out_ranks = [(len(size) + 2,)]
 

nnsmith/abstract/tensor.py

@@ -10,6 +10,9 @@
 
 class AbsTensor:
     def __init__(self, shape: List[Union[int, z3.ExprRef]], dtype: DType):
+        assert isinstance(
+            shape, (list, tuple)
+        ), f"Shape must be a list/tuple, but got {shape}"
         self.shape = list(shape)
         self.dtype = DType(dtype)
 

nnsmith/backends/factory.py

@@ -155,7 +155,9 @@ def make_testcase(
                 log=traceback.format_exc(),
             )
 
-        return TestCase(model, Oracle(input=input, output=output))
+        return TestCase(
+            model, Oracle(input=input, output=output, provider=self.system_name)
+        )
 
     @staticmethod
     def init(name, device="cpu", optmax=True, catch_process_crash=False, **kwargs):
@@ -188,5 +190,23 @@ def init(name, device="cpu", optmax=True, catch_process_crash=False, **kwargs):
                 catch_process_crash=catch_process_crash,
                 **kwargs,
             )
+        elif name == "tflite":
+            from nnsmith.backends.tflite import TFLiteFactory
+
+            return TFLiteFactory(
+                device=device,
+                optmax=optmax,
+                catch_process_crash=catch_process_crash,
+                **kwargs,
+            )
+        elif name == "xla":
+            from nnsmith.backends.xla import XLAFactory
+
+            return XLAFactory(
+                device=device,
+                optmax=optmax,
+                catch_process_crash=catch_process_crash,
+                **kwargs,
+            )
         else:
             raise ValueError(f"unknown backend: {name}")

nnsmith/backends/onnxruntime.py

@@ -22,7 +22,7 @@ def __init__(self, device, optmax, **kwargs):
         super().__init__(device, optmax, **kwargs)
         self.opt_level = OPT_LEVELS[-1 if optmax else 0]
         self.providers = ["CPUExecutionProvider"]
-        if device in ["cuda", "gpu"]:
+        if device == "cuda":
             self.providers = [
                 "CUDAExecutionProvider",
                 "CPUExecutionProvider",

nnsmith/backends/tensorrt.py

@@ -25,16 +25,16 @@ class HostDeviceMem:
 
 
 class TRTFactory(BackendFactory):
-    def __init__(self, device="gpu", optmax=True, **kwargs):
+    def __init__(self, device="cuda", optmax=True, **kwargs):
         super().__init__(device, optmax, **kwargs)
 
-        if device != "gpu":
+        if device != "cuda":
             raise ValueError("TensorRT backend only supports GPU!")
 
         if optmax is False:
             # TODO(@ganler): support non-optimized TensorRT by using performing
             # inference over a model that marks all nodes as outputs.
-            warnings.warn("There is not O0 mode for TensorRT so far.", UserWarning)
+            warnings.warning("There is not O0 mode for TensorRT so far.", UserWarning)
 
     @property
     def system_name(self) -> str:

nnsmith/backends/tflite.py

@@ -43,7 +43,7 @@ def __init__(self, device, optmax, **kwargs) -> None:
 
     @property
     def system_name(self) -> str:
-        "tflite"
+        return "tflite"
 
     @dispatch(TFModel)
     def make_backend(

nnsmith/backends/tvm.py

@@ -24,7 +24,7 @@ def __init__(self, device="cpu", optmax=True, executor="graph", **kwargs) -> Non
         self.opt_level = 4 if optmax else 0
         if device == "cpu":
             self.target = tvm.target.Target("llvm")
-        elif device == "cuda" or device == "gpu":
+        elif device == "cuda":
             self.target = tvm.target.Target("cuda")
         else:
             raise ValueError(f"Unknown device `{device}`")

nnsmith/backends/xla.py

@@ -0,0 +1,46 @@
+from os import PathLike
+from typing import Callable, Dict
+
+import numpy as np
+import tensorflow as tf  # type: ignore
+from multipledispatch import dispatch
+
+from nnsmith.backends.factory import BackendCallable, BackendFactory
+from nnsmith.materialize.tensorflow import (
+    TFModel,
+    TFNetCallable,
+    np_dict_from_tf,
+    tf_dict_from_np,
+)
+
+
+class XLAFactory(BackendFactory):
+    def __init__(self, device="cpu", optmax: bool = False, catch_process_crash=True):
+        super().__init__(device, optmax, catch_process_crash)
+
+    @property
+    def system_name(self) -> str:
+        return "xla"
+
+    @dispatch(TFModel)
+    def make_backend(self, model: TFModel) -> BackendCallable:
+        concrete_net: TFNetCallable = model.concrete_net()
+        device: tf.device
+
+        if self.device == "cpu":
+            device = tf.device(tf.config.list_logical_devices("CPU")[0].name)
+        elif self.device == "cuda":
+            device = tf.device(tf.config.list_logical_devices("GPU")[0].name)
+        else:
+            raise ValueError(f"Unknown device: {self.device}")
+
+        @tf.function(jit_compile=True)
+        def compiled_net(**inputs) -> Dict[str, tf.Tensor]:
+            return concrete_net(**inputs)
+
+        def closure(inputs: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
+            tf.config.run_functions_eagerly(False)
+            with device:
+                return np_dict_from_tf(compiled_net(**tf_dict_from_np(inputs)))
+
+        return closure

nnsmith/cli/dtype_test.py

@@ -9,11 +9,14 @@
 @hydra.main(version_base=None, config_path="../config", config_name="main")
 def main(cfg: DictConfig):
     backend_cfg = cfg["backend"]
-    factory = BackendFactory.init(
-        name=backend_cfg["type"],
-        device=backend_cfg["device"],
-        optmax=backend_cfg["optmax"],
-    )
+    if backend_cfg is not None:
+        factory = BackendFactory.init(
+            name=backend_cfg["type"],
+            device=backend_cfg["device"],
+            optmax=backend_cfg["optmax"],
+        )
+    else:
+        factory = None
     model_type = Model.init(cfg["model"]["type"])
     load_topset_from_auto_cache(model_type, factory)