Update base for Update on "[WIP] Compute forward grads for saved_mean and saved_var when input requires grad"

We want to avoid having to recompute saved_mean and saved_invstd in batch_norm_backward's decomposition in functorch (see pytorch/functorch#877), but also avoid unnecessarily computing forward grads for saved_mean and saved_invstd when they are not needed.

Tested locally with: `python test/test_ops.py -k test_jvpvjp_nn_functional_batch_norm`

Issues:
- not sure if gradgrad in core is missing something, but it is able to pass while the fwgrad_bwgrad comparison fails in functorch


[ghstack-poisoned]

Author: soulitzer

.circleci/cimodel/data/dimensions.py

@@ -4,6 +4,7 @@
     "102",
     "113",
     "116",
+    "117",
 ]
 
 ROCM_VERSIONS = [

.circleci/cimodel/data/simple/ios_definitions.py

@@ -1,4 +1,5 @@
 from cimodel.data.simple.util.versions import MultiPartVersion
+from cimodel.data.simple.util.branch_filters import gen_filter_dict_exclude
 import cimodel.lib.miniutils as miniutils
 
 XCODE_VERSION = MultiPartVersion([12, 5, 1])
@@ -11,7 +12,7 @@ def __init__(self, name, custom_build_name=""):
 
     def render(self):
         extra_parts = [self.custom_build_name] if len(self.custom_build_name) > 0 else []
-        return "_".join([self.name] + extra_parts)
+        return "-".join([self.name] + extra_parts).replace("_", "-")
 
 
 def get_platform(arch_variant_name):
@@ -25,30 +26,25 @@ def __init__(self, xcode_version, arch_variant, is_org_member_context=True, extr
         self.is_org_member_context = is_org_member_context
         self.extra_props = extra_props
 
-    def gen_name_parts(self, with_version_dots):
-
-        version_parts = self.xcode_version.render_dots_or_parts(with_version_dots)
-        build_variant_suffix = "_".join([self.arch_variant.render(), "build"])
-
+    def gen_name_parts(self):
+        version_parts = self.xcode_version.render_dots_or_parts("-")
+        build_variant_suffix = self.arch_variant.render()
         return [
-            "pytorch",
             "ios",
         ] + version_parts + [
             build_variant_suffix,
         ]
 
     def gen_job_name(self):
-        return "_".join(self.gen_name_parts(False))
+        return "-".join(self.gen_name_parts())
 
     def gen_tree(self):
-
         platform_name = get_platform(self.arch_variant.name)
-
         props_dict = {
-            "build_environment": "-".join(self.gen_name_parts(True)),
+            "name": self.gen_job_name(),
+            "build_environment": self.gen_job_name(),
             "ios_arch": self.arch_variant.name,
             "ios_platform": platform_name,
-            "name": self.gen_job_name(),
         }
 
         if self.is_org_member_context:
@@ -57,30 +53,28 @@ def gen_tree(self):
         if self.extra_props:
             props_dict.update(self.extra_props)
 
+        props_dict["filters"] = gen_filter_dict_exclude()
+
         return [{"pytorch_ios_build": props_dict}]
 
 
 WORKFLOW_DATA = [
     IOSJob(XCODE_VERSION, ArchVariant("x86_64"), is_org_member_context=False, extra_props={
         "lite_interpreter": miniutils.quote(str(int(True)))}),
-    IOSJob(XCODE_VERSION, ArchVariant("x86_64", "full_jit"), is_org_member_context=False, extra_props={
-        "lite_interpreter": miniutils.quote(str(int(False)))}),
-    IOSJob(XCODE_VERSION, ArchVariant("arm64"), extra_props={
-        "lite_interpreter": miniutils.quote(str(int(True)))}),
-    IOSJob(XCODE_VERSION, ArchVariant("arm64", "metal"), extra_props={
-        "use_metal": miniutils.quote(str(int(True))),
-        "lite_interpreter": miniutils.quote(str(int(True)))}),
-    IOSJob(XCODE_VERSION, ArchVariant("arm64", "full_jit"), extra_props={
-        "lite_interpreter": miniutils.quote(str(int(False)))}),
-    IOSJob(XCODE_VERSION, ArchVariant("arm64", "custom"), extra_props={
-        "op_list": "mobilenetv2.yaml",
-        "lite_interpreter": miniutils.quote(str(int(True)))}),
+    # IOSJob(XCODE_VERSION, ArchVariant("arm64"), extra_props={
+    #     "lite_interpreter": miniutils.quote(str(int(True)))}),
+    # IOSJob(XCODE_VERSION, ArchVariant("arm64", "metal"), extra_props={
+    #     "use_metal": miniutils.quote(str(int(True))),
+    #     "lite_interpreter": miniutils.quote(str(int(True)))}),
+    # IOSJob(XCODE_VERSION, ArchVariant("arm64", "custom-ops"), extra_props={
+    #     "op_list": "mobilenetv2.yaml",
+    #     "lite_interpreter": miniutils.quote(str(int(True)))}),
     IOSJob(XCODE_VERSION, ArchVariant("x86_64", "coreml"), is_org_member_context=False, extra_props={
         "use_coreml": miniutils.quote(str(int(True))),
         "lite_interpreter": miniutils.quote(str(int(True)))}),
-    IOSJob(XCODE_VERSION, ArchVariant("arm64", "coreml"), extra_props={
-        "use_coreml": miniutils.quote(str(int(True))),
-        "lite_interpreter": miniutils.quote(str(int(True)))}),
+    # IOSJob(XCODE_VERSION, ArchVariant("arm64", "coreml"), extra_props={
+    #     "use_coreml": miniutils.quote(str(int(True))),
+    #     "lite_interpreter": miniutils.quote(str(int(True)))}),
 ]
 
 

.circleci/cimodel/data/simple/macos_definitions.py

@@ -1,3 +1,8 @@
+from collections import OrderedDict
+from cimodel.lib.miniutils import quote
+from cimodel.data.simple.util.branch_filters import gen_filter_dict_exclude
+
+
 class MacOsJob:
     def __init__(self, os_version, is_build=False, is_test=False, extra_props=tuple()):
         # extra_props is tuple type, because mutable data structures for argument defaults
@@ -11,10 +16,14 @@ def gen_tree(self):
         non_phase_parts = ["pytorch", "macos", self.os_version, "py3"]
 
         extra_name_list = [name for name, exist in self.extra_props.items() if exist]
-        full_job_name_list = non_phase_parts + extra_name_list + [
-            'build' if self.is_build else None,
-            'test' if self.is_test else None,
-        ]
+        full_job_name_list = (
+            non_phase_parts
+            + extra_name_list
+            + [
+                "build" if self.is_build else None,
+                "test" if self.is_test else None,
+            ]
+        )
 
         full_job_name = "_".join(list(filter(None, full_job_name_list)))
 
@@ -41,12 +50,99 @@ def gen_tree(self):
         "10_13",
         is_build=True,
         is_test=True,
-        extra_props=tuple({
-            "lite_interpreter": True
-        }.items()),
-    )
+        extra_props=tuple({"lite_interpreter": True}.items()),
+    ),
 ]
 
 
+def get_new_workflow_jobs():
+    return [
+        OrderedDict(
+            {
+                "mac_build": OrderedDict(
+                    {
+                        "name": "macos-12-py3-x86-64-build",
+                        "build-environment": "macos-12-py3-x86-64",
+                        "xcode-version": quote("13.3.1"),
+                        "filters": gen_filter_dict_exclude()
+                    }
+                )
+            }
+        ),
+        OrderedDict(
+            {
+                "mac_test": OrderedDict(
+                    {
+                        "name": "macos-12-py3-x86-64-test-1-2-default",
+                        "build-environment": "macos-12-py3-x86-64",
+                        "xcode-version": quote("13.3.1"),
+                        "shard-number": quote("1"),
+                        "num-test-shards": quote("2"),
+                        "requires": ["macos-12-py3-x86-64-build"],
+                        "filters": gen_filter_dict_exclude()
+                    }
+                )
+            }
+        ),
+        OrderedDict(
+            {
+                "mac_test": OrderedDict(
+                    {
+                        "name": "macos-12-py3-x86-64-test-2-2-default",
+                        "build-environment": "macos-12-py3-x86-64",
+                        "xcode-version": quote("13.3.1"),
+                        "shard-number": quote("2"),
+                        "num-test-shards": quote("2"),
+                        "requires": ["macos-12-py3-x86-64-build"],
+                        "filters": gen_filter_dict_exclude()
+                    }
+                )
+            }
+        ),
+        OrderedDict(
+            {
+                "mac_test": OrderedDict(
+                    {
+                        "name": "macos-12-py3-x86-64-test-1-1-functorch",
+                        "build-environment": "macos-12-py3-x86-64",
+                        "xcode-version": quote("13.3.1"),
+                        "shard-number": quote("1"),
+                        "num-test-shards": quote("1"),
+                        "test-config": "functorch",
+                        "requires": ["macos-12-py3-x86-64-build"],
+                        "filters": gen_filter_dict_exclude()
+                    }
+                )
+            }
+        ),
+        OrderedDict(
+            {
+                "mac_build": OrderedDict(
+                    {
+                        "name": "macos-12-py3-x86-64-lite-interpreter-build-test",
+                        "build-environment": "macos-12-py3-lite-interpreter-x86-64",
+                        "xcode-version": quote("13.3.1"),
+                        "build-generates-artifacts": "false",
+                        "filters": gen_filter_dict_exclude()
+                    }
+                )
+            }
+        ),
+        OrderedDict(
+            {
+                "mac_build": OrderedDict(
+                    {
+                        "name": "macos-12-py3-arm64-build",
+                        "build-environment": "macos-12-py3-arm64",
+                        "xcode-version": quote("13.3.1"),
+                        "python-version": quote("3.9.12"),
+                        "filters": gen_filter_dict_exclude()
+                    }
+                )
+            }
+        ),
+    ]
+
+
 def get_workflow_jobs():
     return [item.gen_tree() for item in WORKFLOW_DATA]

.circleci/cimodel/data/simple/nightly_ios.py

@@ -15,7 +15,7 @@ def __init__(self,
     def get_phase_name(self):
         return "upload" if self.is_upload else "build"
 
-    def get_common_name_pieces(self, with_version_dots):
+    def get_common_name_pieces(self, sep):
 
         extra_name_suffix = [self.get_phase_name()] if self.is_upload else []
 
@@ -24,7 +24,7 @@ def get_common_name_pieces(self, with_version_dots):
         common_name_pieces = [
             "ios",
         ] + extra_name + [
-        ] + ios_definitions.XCODE_VERSION.render_dots_or_parts(with_version_dots) + [
+        ] + ios_definitions.XCODE_VERSION.render_dots_or_parts(sep) + [
             "nightly",
             self.variant,
             "build",
@@ -33,14 +33,14 @@ def get_common_name_pieces(self, with_version_dots):
         return common_name_pieces
 
     def gen_job_name(self):
-        return "_".join(["pytorch"] + self.get_common_name_pieces(False))
+        return "_".join(["pytorch"] + self.get_common_name_pieces(None))
 
     def gen_tree(self):
         build_configs = BUILD_CONFIGS_FULL_JIT if self.is_full_jit else BUILD_CONFIGS
         extra_requires = [x.gen_job_name() for x in build_configs] if self.is_upload else []
 
         props_dict = {
-            "build_environment": "-".join(["libtorch"] + self.get_common_name_pieces(True)),
+            "build_environment": "-".join(["libtorch"] + self.get_common_name_pieces(".")),
             "requires": extra_requires,
             "context": "org-member",
             "filters": {"branches": {"only": "nightly"}},

.circleci/cimodel/data/simple/upload_test_stats_definition.py

@@ -0,0 +1,22 @@
+from typing import OrderedDict
+from cimodel.data.simple.util.branch_filters import gen_filter_dict_exclude
+
+
+def get_workflow_job():
+    return [
+        OrderedDict(
+            {
+                "upload_test_stats": OrderedDict(
+                    {
+                        "name": "upload test status",
+                        "requires": [
+                            "macos-12-py3-x86-64-test-1-2-default",
+                            "macos-12-py3-x86-64-test-2-2-default",
+                            "macos-12-py3-x86-64-test-1-1-functorch",
+                        ],
+                        "filters": gen_filter_dict_exclude()
+                    }
+                )
+            }
+        ),
+    ]

.circleci/cimodel/data/simple/util/branch_filters.py

@@ -12,6 +12,9 @@
 
 RC_PATTERN = r"/v[0-9]+(\.[0-9]+)*-rc[0-9]+/"
 
+MAC_IOS_EXCLUSION_LIST = ["nightly", "postnightly"]
+
+
 def gen_filter_dict(
         branches_list=NON_PR_BRANCH_LIST,
         tags_list=None
@@ -26,3 +29,11 @@ def gen_filter_dict(
     if tags_list is not None:
         filter_dict["tags"] = {"only": tags_list}
     return filter_dict
+
+
+def gen_filter_dict_exclude(branches_list=MAC_IOS_EXCLUSION_LIST):
+    return {
+        "branches": {
+            "ignore": branches_list,
+        },
+    }

.circleci/cimodel/data/simple/util/versions.py

@@ -1,3 +1,6 @@
+from typing import Optional
+
+
 class MultiPartVersion:
     def __init__(self, parts, prefix=""):
         self.parts = parts
@@ -13,14 +16,11 @@ def prefixed_parts(self):
         else:
             return [self.prefix]
 
-    def render_dots(self):
-        return ".".join(self.prefixed_parts())
-
-    def render_dots_or_parts(self, with_dots):
-        if with_dots:
-            return [self.render_dots()]
-        else:
+    def render_dots_or_parts(self, sep: Optional[str] = None):
+        if sep is None:
             return self.prefixed_parts()
+        else:
+            return [sep.join(self.prefixed_parts())]
 
 
 class CudaVersion(MultiPartVersion):