Improving filter for undef behavior and vulops

nnsmith/abstract/op.py

@@ -57,16 +57,18 @@ def __call__(self, op_type: Type["AbsOpBase"]) -> Type["AbsOpBase"]:
 
 
 class mark_materialize:
-    def __init__(self, dialect: str):
+    def __init__(self, dialect: str, limit_domain=False):
+        self.limit_domain = limit_domain
         self.dialect = dialect
 
     def __call__(self, op_type: Type["AbsOpBase"]) -> Type["AbsOpBase"]:
         op_list = FULL_OPERATOR_SETS.setdefault(self.dialect, [])
 
         if op_type not in op_list:
+            op_type = mark_abstract(self.dialect)(op_type)
+            op_type.limit_domain = self.limit_domain
             op_list.append(op_type)
             op_list.sort(key=lambda x: x.__name__)
-            op_type = mark_abstract(self.dialect)(op_type)
 
         return op_type
 
@@ -265,8 +267,8 @@ class AbsOpBase(ABC):
     # this op can accept one of float32xfloat32, float64xfloat64, and int32xint32 as input dtypes.
     in_dtypes: List[Tuple[DType, ...]] = None  # Overwrite me!
     out_dtypes: List[Tuple[DType, ...]] = None
-    # whether to disable the op during graph generation
-    _skip = False
+
+    limit_domain = False
 
     dialect = None
 
@@ -686,7 +688,7 @@ class Div(BcastBinaryOp):
     out_dtypes = [(i,) for i in DTYPE_FLOATS]
 
 
-@mark_materialize("core")
+@mark_materialize("core", limit_domain=True)
 class Pow(BcastBinaryOp):
     in_dtypes = [(i, i) for i in DTYPE_FLOATS]
     out_dtypes = [(i,) for i in DTYPE_FLOATS]
@@ -737,13 +739,13 @@ class Cos(TrigonometricOp):
     out_dtypes = [(i,) for i in DTYPE_FLOATS]
 
 
-@mark_materialize("core")
+@mark_materialize("core", limit_domain=True)
 class Asin(TrigonometricOp):
     in_dtypes = [(i,) for i in DTYPE_FLOATS]
     out_dtypes = [(i,) for i in DTYPE_FLOATS]
 
 
-@mark_materialize("core")
+@mark_materialize("core", limit_domain=True)
 class Acos(TrigonometricOp):
     in_dtypes = [(i,) for i in DTYPE_FLOATS]
     out_dtypes = [(i,) for i in DTYPE_FLOATS]
@@ -795,13 +797,13 @@ class Round(ElementWiseUnaryOp):
     out_dtypes = [(i,) for i in DTYPE_FLOATS]
 
 
-@mark_materialize("core")
+@mark_materialize("core", limit_domain=True)
 class Sqrt(ElementWiseUnaryOp):
     in_dtypes = [(i,) for i in DTYPE_FLOATS]
     out_dtypes = [(i,) for i in DTYPE_FLOATS]
 
 
-@mark_materialize("core")
+@mark_materialize("core", limit_domain=True)
 class Log2(ElementWiseUnaryOp):
     in_dtypes = [(i,) for i in DTYPE_FLOATS]
     out_dtypes = [(i,) for i in DTYPE_FLOATS]

nnsmith/cli/fuzz.py

@@ -129,7 +129,9 @@ def __init__(
             model_cfg["type"], backend_target=cfg["backend"]["target"]
         )
         self.ModelType.add_seed_setter()
-        self.opset = auto_opset(self.ModelType, self.factory)
+        self.opset = auto_opset(
+            self.ModelType, self.factory, vulops=cfg["mgen"]["vulops"]
+        )
 
         seed = cfg["fuzz"]["seed"] or random.getrandbits(32)
         set_seed(seed)

nnsmith/cli/model_gen.py

@@ -38,7 +38,7 @@ def main(cfg: DictConfig):
         factory = None
 
     gen = random_model_gen(
-        opset=auto_opset(ModelType, factory),
+        opset=auto_opset(ModelType, factory, vulops=mgen_cfg["vulops"]),
         init_rank=mgen_cfg["init_rank"],
         seed=seed,
         max_nodes=mgen_cfg["max_nodes"],

nnsmith/config/main.yaml

@@ -29,6 +29,7 @@ mgen: # model gen.
   init_rank: 4
   max_nodes: 5
   timeout_ms: 50000
+  vulops: False
   save: "nnsmith_output"
   seed: null
 

nnsmith/filter.py

@@ -42,7 +42,11 @@ def filter_nan(report: BugReport) -> bool:  # True means filter;
 
     # numpy.assert_allclose style.
     # TODO(ganler): can we use more well-formed checking? say directly checking the results?
-    return "nan location mismatch" in report.log
+    return (
+        "nan location mismatch" in report.log
+        or "-9223372036854775808" in report.log  # tf.cast(nan, int) is UB.
+        or "-2147483648" in report.log
+    )
 
 
 @filter("inf")

nnsmith/narrow_spec.py

@@ -252,12 +252,16 @@ def auto_opconfig(
         return opset
 
 
-def auto_opset(model_cls: Type[Model], factory: Optional[BackendFactory] = None):
+def auto_opset(
+    model_cls: Type[Model],
+    factory: Optional[BackendFactory] = None,
+    vulops: bool = False,
+) -> List[Type[AbsOpBase]]:
     # None means only test model exportation.
     topset_config = auto_opconfig(model_cls, factory)
     opset = []
     for op in model_cls.operators():
-        if op.name() not in topset_config:
+        if op.name() not in topset_config or (vulops == False and op.limit_domain):
             continue
         op.in_dtypes = topset_config[op.name()].in_dtypes
         op.out_dtypes = topset_config[op.name()].out_dtypes