diff --git a/nnsmith/abstract/op.py b/nnsmith/abstract/op.py index a07591ef..e7a39c90 100644 --- a/nnsmith/abstract/op.py +++ b/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] diff --git a/nnsmith/cli/fuzz.py b/nnsmith/cli/fuzz.py index f38cedda..2d83d256 100644 --- a/nnsmith/cli/fuzz.py +++ b/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) diff --git a/nnsmith/cli/model_gen.py b/nnsmith/cli/model_gen.py index 4bbfdc67..fefd4fc7 100644 --- a/nnsmith/cli/model_gen.py +++ b/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"], diff --git a/nnsmith/config/main.yaml b/nnsmith/config/main.yaml index 7c287b63..f43f950c 100644 --- a/nnsmith/config/main.yaml +++ b/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 diff --git a/nnsmith/filter.py b/nnsmith/filter.py index 24733b1a..5aeafd3c 100644 --- a/nnsmith/filter.py +++ b/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") diff --git a/nnsmith/narrow_spec.py b/nnsmith/narrow_spec.py index 1f98d6c3..a5e3e8a9 100644 --- a/nnsmith/narrow_spec.py +++ b/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