ci

Author: xadupre

_unittests/ut_helpers/test_doc_helper.py

@@ -3,15 +3,15 @@
 import onnx.helper as oh
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase
-from onnx_diagnostic.helpers.doc_helper import LayerNormalizationOrt
+from onnx_diagnostic.helpers.doc_helper import LayerNormalizationOrt, MatMulOrt
 from onnx_diagnostic.reference import TorchOnnxEvaluator
 
 TFLOAT = onnx.TensorProto.FLOAT
 TFLOAT16 = onnx.TensorProto.FLOAT16
 
 
 class TestDocHelper(ExtTestCase):
-    def test_custom_doc_kernels(self):
+    def test_custom_doc_kernels_layer_normalization(self):
         model = oh.make_model(
             oh.make_graph(
                 [
@@ -35,7 +35,7 @@ def test_custom_doc_kernels(self):
                 [oh.make_tensor_value_info("Z", TFLOAT16, ["b", "c", "d"])],
             ),
             ir_version=9,
-            opset_imports=[oh.make_opsetid("", 17)],
+            opset_imports=[oh.make_opsetid("", 18)],
         )
 
         torch_sess = TorchOnnxEvaluator(model, verbose=0)
@@ -58,6 +58,40 @@ def test_custom_doc_kernels(self):
         got = torch_sess_custom.run(None, feeds)
         self.assertEqualAny(expected, got)
 
+    def test_custom_doc_kernels_matmul(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("MatMul", ["X", "Y"], ["Z"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT16, ["b", "c", "d"]),
+                    oh.make_tensor_value_info("Y", TFLOAT16, ["b", "d", "e"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT16, ["b", "c", "e"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+
+        torch_sess = TorchOnnxEvaluator(model, verbose=0)
+        torch_sess_custom = TorchOnnxEvaluator(
+            model,
+            verbose=0,
+            custom_kernels={("", "MatMul"): MatMulOrt},
+        )
+        feeds = dict(
+            zip(
+                torch_sess.input_names,
+                [
+                    torch.rand(3, 4, 5, dtype=torch.float16),
+                    torch.rand(3, 5, 7, dtype=torch.float16),
+                ],
+            )
+        )
+        expected = torch_sess.run(None, feeds)
+        got = torch_sess_custom.run(None, feeds)
+        self.assertEqualAny(expected, got)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/helpers/doc_helper.py

@@ -7,7 +7,7 @@
 
 
 class LayerNormalizationOrt(OpRunKernel):
-    "LayerNormalization"
+    "LayerNormalization with onnxruntime"
 
     @classmethod
     def device_dependent(cls) -> bool:
@@ -25,7 +25,7 @@ def __init__(
         self.axis = self.get_attribute_int(node, "axis", -1)
         self.epsilon = self.get_attribute_float(node, "epsilon", 1e-5)
         self.device = device
-        self.stash_type = onnx_dtype_to_torch_dtype(
+        self.stash_type = onnx_dtype_to_torch_dtype(  # type: ignore[arg-type]
             self.get_attribute_int(node, "stash_type", onnx.TensorProto.FLOAT)
         )
         self.compute_std = len(node.output) > 1
@@ -36,7 +36,7 @@ def __init__(
         self._cache: Dict[Tuple[int, int], onnx.ModelProto] = {}
         self.is_cpu = torch.device("cpu") == self.device
 
-    def _make_model(self, dtype: int, rank: int) -> onnx.ModelProto:
+    def _make_model(self, itype: int, rank: int) -> onnx.ModelProto:
         shape = [*["d{i}" for i in range(rank - 1)], "last"]
         layer_model = oh.make_model(
             oh.make_graph(
@@ -51,14 +51,14 @@ def _make_model(self, dtype: int, rank: int) -> onnx.ModelProto:
                 ],
                 "dummy",
                 [
-                    oh.make_tensor_value_info("X", onnx.TensorProto.FLOAT16, shape),
-                    oh.make_tensor_value_info("W", onnx.TensorProto.FLOAT16, ["last"]),
-                    oh.make_tensor_value_info("B", onnx.TensorProto.FLOAT16, ["last"]),
+                    oh.make_tensor_value_info("X", itype, shape),
+                    oh.make_tensor_value_info("W", itype, ["last"]),
+                    oh.make_tensor_value_info("B", itype, ["last"]),
                 ],
-                [oh.make_tensor_value_info("Z", onnx.TensorProto.FLOAT16, shape)],
+                [oh.make_tensor_value_info("Z", itype, shape)],
             ),
             ir_version=9,
-            opset_imports=[oh.make_opsetid("", 17)],
+            opset_imports=[oh.make_opsetid("", 18)],
         )
         import onnxruntime
 
@@ -80,3 +80,58 @@ def run(self, x, scale, bias=None):
         feeds = {k: v.tensor.detach().cpu().numpy() for k, v in feeds.items()}
         got = sess.run(None, feeds)[0]
         return OpRunTensor(torch.from_numpy(got).to(x.dtype).to(x.device))
+
+
+class MatMulOrt(OpRunKernel):
+    "MatMul with onnxruntime"
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        "Needs device."
+        return False
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version=None,
+        device: Optional[torch.device] = None,
+        verbose=0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.device = device
+        self._cache: Dict[Tuple[int, int, int], onnx.ModelProto] = {}
+        self.is_cpu = torch.device("cpu") == self.device
+
+    def _make_model(self, itype: int, ranka: int, rankb: int) -> onnx.ModelProto:
+        shapea = ["a{i}" for i in range(ranka)]
+        shapeb = ["b{i}" for i in range(rankb)]
+        shapec = ["c{i}" for i in range(max(ranka, rankb))]
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("MatMul", ["A", "B"], ["C"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("A", itype, shapea),
+                    oh.make_tensor_value_info("B", itype, shapeb),
+                ],
+                [oh.make_tensor_value_info("C", itype, shapec)],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 17)],
+        )
+        import onnxruntime
+
+        provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
+        return onnxruntime.InferenceSession(model.SerializeToString(), providers=[provider])
+
+    def run(self, a, b):
+        itype = torch_dtype_to_onnx_dtype(a.dtype)
+        ranka, rankb = len(a.shape), len(b.shape)
+        key = itype, ranka, rankb
+        if key not in self._cache:
+            self._cache[key] = self._make_model(itype, ranka, rankb)
+        sess = self._cache[key]
+        feeds = dict(A=a, B=b)
+        feeds = {k: v.tensor.detach().cpu().numpy() for k, v in feeds.items()}
+        got = sess.run(None, feeds)[0]
+        return OpRunTensor(torch.from_numpy(got).to(a.dtype).to(a.device))