Calibrator: devices

Reviewed By: SS-JIA

Differential Revision: D54544308

fbshipit-source-id: 8b7ef524786d566b869ff21d2b14e64690400514

Author: manuelcandales

calibrator/runner.py

@@ -9,6 +9,7 @@
 import sys
 from typing import Any, List, OrderedDict, Tuple
 
+import torch
 from executorch.exir.dialects.edge.op.api import get_callable, to_variant
 from inputgen.argtuple.engine import MetaArgTupleEngine
 from inputgen.argtuple.gen import ArgumentTupleGenerator
@@ -26,13 +27,25 @@ def smt(meta_tuple):
 
 
 class SpecRunner:
-    def __init__(self, spec: Spec, *, valid: bool = True, out: bool = False):
+    def __init__(
+        self,
+        spec: Spec,
+        *,
+        valid: bool = True,
+        out: bool = False,
+        devices: Tuple[str] = ("cpu",),
+    ):
         self.spec = spec
         self.generator = ArgumentTupleGenerator(self.spec)
         self.valid = valid
         self.out = out
         self.op_name = spec.op
         self.op = self.get_callable_op()
+        self.results = {}
+        self.devices = devices
+        self.results = {}
+        for device in self.devices:
+            self.results[device] = {}
 
     def get_callable_op(self):
         name = self.spec.op
@@ -42,26 +55,93 @@ def get_callable_op(self):
             op: OpOverload = to_variant(op, SchemaKind.out)
         return op
 
-    def run(self):
+    def report_device(self, device):
+        print(f"Device: {device}\n")
         failures = []
-        engine = MetaArgTupleEngine(self.spec, out=self.out)
-        for meta_tuple in engine.gen(valid=self.valid):
-            success, _, _, _, _ = self.run_meta_tuple(meta_tuple)
+        for meta_tuple in self.results[device]:
+            success = self.results[device][meta_tuple]
             if not success:
                 failures.append(meta_tuple)
         if len(failures) > 0:
             print("FAILURES\n")
             for meta_tuple in failures:
-                print(f"\t{smt(meta_tuple)}\n")
+                print(f"\t{meta_tuple}\n")
         else:
             print("SUCCESS\n")
 
+    def report_inconsistencies(self):
+        print(f"Devices: {' '.join(self.devices)}\n")
+        meta_tuples = self.results[self.devices[0]].keys()
+        inconsistencies = set()
+        for meta_tuple in meta_tuples:
+            res = self.results[self.devices[0]][meta_tuple]
+            for device in self.devices[1:]:
+                res ^= self.results[device][meta_tuple]
+                if not res:
+                    inconsistencies.append(meta_tuple)
+        if len(inconsistencies) > 0:
+            print("INCONSISTENCIES\n")
+            for meta_tuple in inconsistencies:
+                res = [self.results[d][meta_tuple] for d in self.devices]
+                res_string = " ".join(["x" if r else "o" for r in res])
+                print(f"\t{res_string} {meta_tuple}\n")
+
+    def run(self):
+        engine = MetaArgTupleEngine(self.spec, out=self.out)
+        for meta_tuple in engine.gen(valid=self.valid):
+            self.run_meta_tuple(meta_tuple)
+        if len(self.devices) > 1:
+            self.report_inconsistencies()
+        for device in self.devices:
+            self.report_device(device)
+
+    def move_to_device(
+        self,
+        device: str,
+        cpu_posargs: List[Any],
+        cpu_inkwargs: OrderedDict[str, Any],
+        cpu_outargs: OrderedDict[str, Any],
+    ):
+        if device == "cpu":
+            return cpu_posargs, cpu_inkwargs, cpu_outargs
+        posargs = []
+        inkwargs = OrderedDict()
+        outargs = OrderedDict()
+        for arg in cpu_posargs:
+            new = arg
+            if isinstance(arg, torch.Tensor):
+                new = arg.to(device=device)
+            posargs.append(new)
+        for k, v in cpu_inkwargs.items():
+            new = v
+            if isinstance(v, torch.Tensor):
+                new = v.to(device=device)
+            inkwargs[k] = new
+        for k, v in cpu_outargs.items():
+            new = v
+            if isinstance(v, torch.Tensor):
+                new = v.to(device=device)
+            outargs[k] = new
+        return posargs, inkwargs, outargs
+
     def run_meta_tuple(
         self, meta_tuple: Tuple[MetaArg]
     ) -> Tuple[bool, Any, List[Any], OrderedDict[str, Any], OrderedDict[str, Any]]:
         print(f"Running op: {self.op_name}, meta_tuple: {[str(x) for x in meta_tuple]}")
         posargs, inkwargs, outargs = self.generator.gen_tuple(meta_tuple, out=self.out)
-        return self.run_values(meta_tuple, posargs, inkwargs, outargs)
+        for device in self.devices:
+            posargs, inkwargs, outargs = self.move_to_device(
+                device, posargs, inkwargs, outargs
+            )
+            success, res, posargs, inkwargs, outargs = self.run_values(
+                meta_tuple, posargs, inkwargs, outargs
+            )
+            mt = smt(meta_tuple)
+            if mt in self.results[device]:
+                logging.warning(f"Repeated meta_tuple {mt}")
+                self.results[device][mt] &= success
+            else:
+                self.results[device][mt] = success
 
     def run_values(
         self,
@@ -96,13 +176,14 @@ def main():
         "--invalid", action="store_true", help="generate invalid inputs"
     )
     parser.add_argument("--out", action="store_true", help="run out variants")
+    parser.add_argument("--devices", nargs="*", default=("cpu",), help="run on devices")
     args = parser.parse_args()
 
     if args.op not in SpecDictDB:
         raise RuntimeError(f"Op {args.op} not found in SpecDB")
 
     spec = SpecDictDB[args.op]
-    SpecRunner(spec, valid=not args.invalid, out=args.out).run()
+    SpecRunner(spec, valid=not args.invalid, out=args.out, devices=args.devices).run()
 
 
 if __name__ == "__main__":

inputgen/argtuple/gen.py

@@ -18,7 +18,7 @@ def __init__(self, spec: Spec):
 
     def gen_tuple(
         self, meta_tuple: Tuple[MetaArg], *, out: bool = False
-    ) -> Tuple[List[Any], OrderedDict[str, Any]]:
+    ) -> Tuple[List[Any], OrderedDict[str, Any], OrderedDict[str, Any]]:
         posargs = []
         inkwargs = OrderedDict()
         outargs = OrderedDict()

specdb/db.py

@@ -997,9 +997,10 @@
                 name="bias",
                 deps=[1, 6, 8],
                 constraints=[
-                    cp.Dtype.In(
-                        lambda deps: [deps[0].dtype] if deps[1] else dt._floating
-                    ),  # TODO(mcandales): Calibrate.
+                    cp.Dtype.Eq(lambda deps: deps[0].dtype),
+                    # cp.Dtype.In(
+                    #    lambda deps: [deps[0].dtype] if deps[1] else dt._floating
+                    # ),  # TODO(mcandales): Calibrate.
                     cp.Rank.Eq(lambda deps: 1),
                     cp.Size.Eq(
                         lambda deps, r, d: fn.conv_bias_size_eq(
@@ -1039,7 +1040,7 @@
                 ArgType.Bool,
                 name="transposed",
                 # TODO(mcandales): Executorch specific constraint
-                # constraints=[cp.Value.In(lambda deps: [False]]),
+                # constraints=[cp.Value.Eq(lambda deps: False)],
             ),
             InPosArg(  # output_padding
                 ArgType.LengthList,