Update base for Update on "add more export modules after ertrecod created"

we need to support etrecord recording custom export modules for further usage. This diff makes that happen by creating new function inside ETRecord

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

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"
     }
-}
+}