Update on "make to_edge_transform_and_lower support etrecord generation"

Differential Revision: [D79336982](https://our.internmc.facebook.com/intern/diff/D79336982/)

umbrella issue: #12961

[ghstack-poisoned]

Author: Gasoonjia

backends/apple/coreml/partition/coreml_partitioner.py

@@ -20,6 +20,7 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
+from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -56,6 +57,80 @@ def log_once(self, msg: str) -> None:
             logger.info(msg)
             self._logged_msgs.add(msg)
 
+    def should_skip_op_for_delegation(self, node_target_name: str) -> bool:
+        skipped_ops = self.skip_ops_for_coreml_delegation or []
+        if node_target_name in skipped_ops:
+            assert (
+                not self.lower_full_graph
+            ), f"Cannot skip {node_target_name} because lower_full_graph is True.  Please set skip_ops_for_coreml_delegation=None or lower_full_graph=False in the CoreMLPartitioner"
+            self.log_once(
+                "Skipping op for CoreML delegation because it is in skip_ops_for_coreml_delegation: "
+                + node_target_name
+            )
+            return True
+        return False
+
+    def should_override_support(self, node) -> bool:
+        # https://github.com/apple/coremltools/issues/2573
+        if (
+            node.target
+            in [
+                torch.ops.aten.sub.Tensor,
+                exir_ops.edge.aten.sub.Tensor,
+                torch.ops.aten.add.Tensor,
+                exir_ops.edge.aten.add.Tensor,
+            ]
+            and "alpha" in node.kwargs
+            and node.kwargs["alpha"] != 1
+        ):
+            self.log_once(
+                "torch.ops.aten.{sub, add}.Tensor with alpha != 1 is not supported by CoreML.  Overriding support."
+            )
+            return True
+
+        # https://github.com/apple/coremltools/issues/2565
+        if node.target in [
+            torch.ops.aten.diagonal.default,
+            torch.ops.aten.diagonal_copy.default,
+            exir_ops.edge.aten.diagonal.default,
+            exir_ops.edge.aten.diagonal_copy.default,
+        ]:
+            self.log_once(
+                "torch.ops.aten.diagonal.default has a bug in CoreML.  Overriding op support."
+            )
+            return True
+
+        # https://github.com/apple/coremltools/issues/2569
+        if node.target in [
+            torch.ops.aten.acosh.default,
+            exir_ops.edge.aten.acosh.default,
+            torch.ops.aten.asinh.default,
+            exir_ops.edge.aten.asinh.default,
+        ]:
+            self.log_once(
+                "torch.ops.aten.{acosh, asinh}.default is not supported by CoreML.  Overriding op support."
+            )
+            return True
+
+        # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
+        # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
+        # # in the placeholders due to partitioning, which CoreML does not support
+        # if not self.lower_full_graph and any(
+        #     isinstance(arg, torch.fx.Node)
+        #     and isinstance(
+        #         arg.meta.get("val", None),
+        #         (torch.SymInt, torch.SymBool, torch.SymFloat),
+        #     )
+        #     for arg in node.args
+        # ):
+        #     self.log_once(
+        #         "Skipping op for CoreML delegation because it contains symbolic args: "
+        #         + node_target_name
+        #     )
+        #     return True
+
+        return False
+
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         # get_attr node can always be supported on any backend
         if node.op == "get_attr":
@@ -64,38 +139,17 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         elif node.op == "call_function":
             # skip ops if specified by user
             node_target_name = getattr(node.target, "__name__", "").lower()
-            if node_target_name in (self.skip_ops_for_coreml_delegation or []):
-                self.log_once(
-                    "Skipping op for CoreML delegation because it is in skip_ops_for_coreml_delegation: "
-                    + node_target_name
-                )
-                assert (
-                    not self.lower_full_graph
-                ), "Cannot have skip_ops_for_coreml_delegation when lower_full_graph is True"
-                return False
 
-            # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
-            # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
-            # # in the placeholders due to partitioning, which CoreML does not support
-            # if not self.lower_full_graph and any(
-            #     isinstance(arg, torch.fx.Node)
-            #     and isinstance(
-            #         arg.meta.get("val", None),
-            #         (torch.SymInt, torch.SymBool, torch.SymFloat),
-            #     )
-            #     for arg in node.args
-            # ):
-            #     self.log_once(
-            #         "Skipping op for CoreML delegation because it contains symbolic args: "
-            #         + node_target_name
-            #     )
-            #     assert not self.lower_full_graph
-            #     return False
+            if self.should_skip_op_for_delegation(node_target_name):
+                return False
 
             # query coremltools to see if node is supported
             is_supported = ct.converters.mil.frontend.torch.is_torch_fx_node_supported(
                 node
             )
+            if self.should_override_support(node):
+                is_supported = False
+
             if not is_supported:
                 if self.lower_full_graph:
                     raise NotImplementedError(
@@ -126,7 +180,6 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 
 
 class CoreMLPartitioner(Partitioner):
-
     def __init__(
         self,
         *,

backends/arm/test/ops/test_asinh.py

@@ -9,10 +9,10 @@
 
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
-    EthosU55PipelineBI,
-    EthosU85PipelineBI,
-    TosaPipelineBI,
-    TosaPipelineMI,
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineFP,
+    TosaPipelineINT,
 )
 
 input_t = Tuple[torch.Tensor]  # Input x
@@ -36,8 +36,8 @@ def forward(self, x):
 
 
 @common.parametrize("test_data", test_data_suite)
-def test_asin_tosa_MI(test_data: Tuple):
-    pipeline = TosaPipelineMI[input_t](
+def test_asinh_tosa_FP(test_data: Tuple):
+    pipeline = TosaPipelineFP[input_t](
         Asinh(),
         (test_data(),),
         aten_op,
@@ -47,8 +47,8 @@ def test_asin_tosa_MI(test_data: Tuple):
 
 
 @common.parametrize("test_data", test_data_suite)
-def test_asin_tosa_BI(test_data: Tuple):
-    pipeline = TosaPipelineBI[input_t](
+def test_asinh_tosa_INT(test_data: Tuple):
+    pipeline = TosaPipelineINT[input_t](
         Asinh(),
         (test_data(),),
         aten_op=[],
@@ -59,8 +59,8 @@ def test_asin_tosa_BI(test_data: Tuple):
 
 @common.parametrize("test_data", test_data_suite)
 @common.XfailIfNoCorstone300
-def test_asin_u55_BI(test_data: Tuple):
-    pipeline = EthosU55PipelineBI[input_t](
+def test_asinh_u55_INT(test_data: Tuple):
+    pipeline = EthosU55PipelineINT[input_t](
         Asinh(),
         (test_data(),),
         aten_ops=[],
@@ -70,8 +70,8 @@ def test_asin_u55_BI(test_data: Tuple):
 
 @common.parametrize("test_data", test_data_suite)
 @common.XfailIfNoCorstone320
-def test_asin_u85_BI(test_data: Tuple):
-    pipeline = EthosU85PipelineBI[input_t](
+def test_asinh_u85_INT(test_data: Tuple):
+    pipeline = EthosU85PipelineINT[input_t](
         Asinh(),
         (test_data(),),
         aten_ops=[],

extension/android/BUCK

@@ -13,9 +13,9 @@ non_fbcode_target(_kind = fb_android_library,
         "executorch_android/src/main/java/org/pytorch/executorch/MethodMetadata.java",
         "executorch_android/src/main/java/org/pytorch/executorch/Module.java",
         "executorch_android/src/main/java/org/pytorch/executorch/Tensor.java",
-        "executorch_android/src/main/java/org/pytorch/executorch/TrainingModule.java",
-        "executorch_android/src/main/java/org/pytorch/executorch/SGD.java",
         "executorch_android/src/main/java/org/pytorch/executorch/annotations/Experimental.java",
+        "executorch_android/src/main/java/org/pytorch/executorch/training/TrainingModule.java",
+        "executorch_android/src/main/java/org/pytorch/executorch/training/SGD.java",
     ],
     autoglob = False,
     language = "JAVA",

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/TrainingModuleE2ETest.kt renamed to extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/training/TrainingModuleE2ETest.kt

@@ -5,21 +5,24 @@
  * This source code is licensed under the BSD-style license found in the
  * LICENSE file in the root directory of this source tree.
  */
-package org.pytorch.executorch
+
+package org.pytorch.executorch.training
 
 import android.Manifest
 import android.util.Log
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.rule.GrantPermissionRule
-import java.io.File
-import java.io.IOException
-import java.net.URISyntaxException
 import org.apache.commons.io.FileUtils
 import org.junit.Assert
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.pytorch.executorch.TestFileUtils.getTestFilePath
+import org.pytorch.executorch.EValue
+import org.pytorch.executorch.Tensor
+import org.pytorch.executorch.TestFileUtils
+import java.io.File
+import java.io.IOException
+import java.net.URISyntaxException
 import kotlin.random.Random
 import kotlin.test.assertContains
 
@@ -36,17 +39,20 @@ class TrainingModuleE2ETest {
         val pteFilePath = "/xor.pte"
         val ptdFilePath = "/xor.ptd"
 
-        val pteFile = File(getTestFilePath(pteFilePath))
+        val pteFile = File(TestFileUtils.getTestFilePath(pteFilePath))
         val pteInputStream = javaClass.getResourceAsStream(pteFilePath)
         FileUtils.copyInputStreamToFile(pteInputStream, pteFile)
         pteInputStream.close()
 
-        val ptdFile = File(getTestFilePath(ptdFilePath))
+        val ptdFile = File(TestFileUtils.getTestFilePath(ptdFilePath))
         val ptdInputStream = javaClass.getResourceAsStream(ptdFilePath)
         FileUtils.copyInputStreamToFile(ptdInputStream, ptdFile)
         ptdInputStream.close()
 
-        val module = TrainingModule.load(getTestFilePath(pteFilePath), getTestFilePath(ptdFilePath))
+        val module = TrainingModule.load(
+            TestFileUtils.getTestFilePath(pteFilePath),
+            TestFileUtils.getTestFilePath(ptdFilePath)
+        )
         val params = module.namedParameters("forward")
 
         Assert.assertEquals(4, params.size)
@@ -75,7 +81,10 @@ class TrainingModuleE2ETest {
             val targetDex = inputDex + 1
             val input = dataset.get(inputDex)
             val target = dataset.get(targetDex)
-            val out = module.executeForwardBackward("forward", EValue.from(input), EValue.from(target))
+            val out = module.executeForwardBackward("forward",
+                EValue.from(input),
+                EValue.from(target)
+            )
             val gradients = module.namedGradients("forward")
 
             if (i == 0) {
@@ -96,7 +105,9 @@ class TrainingModuleE2ETest {
                         input.getDataAsFloatArray()[0],
                         input.getDataAsFloatArray()[1],
                         out[1].toTensor().getDataAsLongArray()[0],
-                        target.getDataAsLongArray()[0]));
+                        target.getDataAsLongArray()[0]
+                    )
+                );
             }
 
             sgd.step(gradients)
@@ -113,12 +124,12 @@ class TrainingModuleE2ETest {
     fun testTrainXOR_PTEOnly() {
         val pteFilePath = "/xor_full.pte"
 
-        val pteFile = File(getTestFilePath(pteFilePath))
+        val pteFile = File(TestFileUtils.getTestFilePath(pteFilePath))
         val pteInputStream = javaClass.getResourceAsStream(pteFilePath)
         FileUtils.copyInputStreamToFile(pteInputStream, pteFile)
         pteInputStream.close()
 
-        val module = TrainingModule.load(getTestFilePath(pteFilePath));
+        val module = TrainingModule.load(TestFileUtils.getTestFilePath(pteFilePath));
         val params = module.namedParameters("forward")
 
         Assert.assertEquals(4, params.size)
@@ -147,7 +158,10 @@ class TrainingModuleE2ETest {
             val targetDex = inputDex + 1
             val input = dataset.get(inputDex)
             val target = dataset.get(targetDex)
-            val out = module.executeForwardBackward("forward", EValue.from(input), EValue.from(target))
+            val out = module.executeForwardBackward("forward",
+                EValue.from(input),
+                EValue.from(target)
+            )
             val gradients = module.namedGradients("forward")
 
             if (i == 0) {
@@ -168,7 +182,9 @@ class TrainingModuleE2ETest {
                         input.getDataAsFloatArray()[0],
                         input.getDataAsFloatArray()[1],
                         out[1].toTensor().getDataAsLongArray()[0],
-                        target.getDataAsLongArray()[0]));
+                        target.getDataAsLongArray()[0]
+                    )
+                );
             }
 
             sgd.step(gradients)
@@ -184,24 +200,33 @@ class TrainingModuleE2ETest {
     @Throws(IOException::class)
     fun testMissingPteFile() {
         val exception = Assert.assertThrows(RuntimeException::class.java) {
-            TrainingModule.load(getTestFilePath(MISSING_PTE_NAME))
+            TrainingModule.load(TestFileUtils.getTestFilePath(MISSING_PTE_NAME))
         }
-        Assert.assertEquals(exception.message, "Cannot load model path!! " + getTestFilePath(MISSING_PTE_NAME))
+        Assert.assertEquals(
+            exception.message,
+            "Cannot load model path!! " + TestFileUtils.getTestFilePath(MISSING_PTE_NAME)
+        )
     }
 
     @Test
     @Throws(IOException::class)
     fun testMissingPtdFile() {
         val exception = Assert.assertThrows(RuntimeException::class.java) {
             val pteFilePath = "/xor.pte"
-            val pteFile = File(getTestFilePath(pteFilePath))
+            val pteFile = File(TestFileUtils.getTestFilePath(pteFilePath))
             val pteInputStream = javaClass.getResourceAsStream(pteFilePath)
             FileUtils.copyInputStreamToFile(pteInputStream, pteFile)
             pteInputStream.close()
 
-            TrainingModule.load(getTestFilePath(pteFilePath), getTestFilePath(MISSING_PTD_NAME))
+            TrainingModule.load(
+                TestFileUtils.getTestFilePath(pteFilePath),
+                TestFileUtils.getTestFilePath(MISSING_PTD_NAME)
+            )
         }
-        Assert.assertEquals(exception.message, "Cannot load data path!! " + getTestFilePath(MISSING_PTD_NAME))
+        Assert.assertEquals(
+            exception.message,
+            "Cannot load data path!! " + TestFileUtils.getTestFilePath(MISSING_PTD_NAME)
+        )
     }
 
     companion object {
@@ -212,4 +237,4 @@ class TrainingModuleE2ETest {
         private const val MISSING_PTE_NAME = "/missing.pte"
         private const val MISSING_PTD_NAME = "/missing.ptd"
     }
-}
+}