Update verifier to handle None Tensor outputs (#8235)

Summary:

conv returns (None, Tensor, Tensor) which is uncommon to see since the schema is (Tensor, Tensor, Tensor). This is to test that the verifier just ignores the None return value (since itll be unused in the runtime).

Reviewed By: larryliu0820

Differential Revision: D69209059

Author: JacobSzwejbka

exir/verification/arg_validator.py

@@ -108,12 +108,23 @@ def call_function(  # noqa: C901  # pyre-fixme[14]
         for schema_ret in target._schema.returns:
             name = schema_ret.name if schema_ret.name else f"__ret_{ret_index}"
             kernel_ret = next(ret_iter)
-            # Return value should not be in OptionalTensor type, so only check torch.TensorType here.
-            if isinstance(schema_ret.type, torch.TensorType) and isinstance(
-                kernel_ret, torch.Tensor
-            ):
-                tensor_arg_types[name] = kernel_ret.dtype
-                ret_index += 1
+            if isinstance(schema_ret.type, torch.TensorType):
+                if isinstance(kernel_ret, torch.Tensor):
+                    tensor_arg_types[name] = kernel_ret.dtype
+                    ret_index += 1
+                # Exceptionally rarely (basically only backwards ops) you might see an OptionalTensor returned.
+                # The schema of these ops though is typically -> (Tensor, Tensor ...). So the actual type
+                # returned in cpp is empty/undefined tensor. There is no analogy to this in python so it
+                # gets crudely mapped to None. To properly fix this core pytorch would have to change the
+                # schema to (Tensor?, ...) which is just never going to happen. So we have to handle this case
+                # here in the verifier and in memory planning as well.
+                elif kernel_ret is None:
+                    tensor_arg_types[name] = schema_ret.default_value
+                    ret_index += 1
+                else:
+                    raise InternalError(
+                        f"encountered return with type Tensor but value wasnt a tensor or None. schema:{target._schema}, output:{ret_index}"
+                    )
             elif schema_ret.type == torch.ListType.ofTensors() and all(
                 isinstance(kernel_ret[i], torch.Tensor) for i in range(len(kernel_ret))
             ):

exir/verification/test/test_verifier.py

@@ -14,9 +14,11 @@
 from executorch.exir import EdgeCompileConfig, to_edge
 
 from executorch.exir.dialects._ops import ops
+from torch import nn
 from torch._export.verifier import SpecViolationError
 from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib  # noqa: F401
 from torch.export import export
+from torch.export.experimental import _export_forward_backward
 
 from ..verifier import EXIREdgeDialectVerifier
 
@@ -123,3 +125,39 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             dim_order_verifier(stride_edge_model.exported_program())
         with self.assertRaises(SpecViolationError):
             stride_verifier(dim_order_edge_model.exported_program())
+
+    def test_none_return_verifier(self) -> None:
+        class Net(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv1 = nn.Conv2d(6, 6, 5)
+                self.linear = nn.Linear(6, 2)
+
+            def forward(self, x):
+                return self.linear(self.conv1(x).flatten(1))
+
+        class TrainingNet(nn.Module):
+            def __init__(self, net):
+                super().__init__()
+                self.net = net
+                self.loss = nn.CrossEntropyLoss()
+
+            def forward(self, input, label):
+                pred = self.net(input)
+                return self.loss(pred, label)
+
+        # conv returns (None, Tensor, Tensor) which is uncommon to see since
+        # the schema is (Tensor, Tensor, Tensor). This is to test that
+        # the verifier just ignores the None return value (since itll be
+        # unused in the runtime).
+        net = TrainingNet(Net())
+        inputs = (torch.randn(1, 6, 5, 5), torch.ones(1, dtype=torch.int64))
+
+        export_model = export(net, inputs)
+        export_model = _export_forward_backward(export_model)
+
+        edge = to_edge(export_model)
+
+        edge_verifier = EXIREdgeDialectVerifier()
+
+        edge_verifier(edge.exported_program())