Update verifier to handle None Tensor outputs

Differential Revision: D69209059

Pull Request resolved: #8235

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())