Update on "[ONNX] Introduce FX-ONNX dispatcher"

Needs microsoft/onnxscript#721

The current FX exporter is using manually maintained dictionary to map ATen op to its OnnxFunction. However, the issue arises when ATen op has overloads or OnnxFunction has overloads, which is not resolvable by the one to one mapping . For example, `aten::arange` has onverloads: `aten::arange.start` and `aten::arange.start_step`, or for `aten::argmax`, torchlib provides two function: aten_argmax, and aten_argmax_dim.

This PR utilizes newly introduced [ONNX OpSchema](microsoft/onnxscript#626) to match the input arguments of an ATen operator to find the correct overload.

### OnnxRegistry

Heavily reference on [TorchScript Registry](#84382). The only difference is that in FX registry, an ATen operator with specific opset version is mapped to a list of overloaded functions.

* No longer use global registry. The registry is initialized in `ResolvedExportOptions` with torchlib, and will be exposed to users in the future.
* Multiple opset version layer is kept through `_SymbolicFunctionGroup` , but torchlib now only supports 18.
* Basic API of custom operator support: `register`, `unregister`, and `is_register_op` are kept for future development. To further complete them, the follow-up PRs should address:
    - How to allow users to remove/override specific overload? Using OpSchema to differentiate?
    - User registers a new overload with the same OpSchema as one of registered overload. 
 
### OnnxDispatcher

Dispatch ATen operators to the matched overload by comparing OpSchema with input arguments.

* `OpSchemaWrapper` wrap the onnx schema, and record matching score.
* `dispatch` uses `OpSchemaWrapper` to compare data types to find the best matched overload. If the match isn't perfect, record warning in diagnostics.
* `dispatch_opset_version` is referenced from #84382 and kept, but torchlib doesn't support opset version != 18.
* Because right now (1) OnnxFunction arguments are manually typed, and (2) ORT could unfollow ONNX type spec, we relax the schema match with `matching score system`.
* To include more supports:  the follow-up PRs should address:
    - How to add op.Cast with autocast? In torchlib or converter?
    - The need of type promotion can be captured by dispatcher, but needs OpSchema shows the T1/T2 information.

### OpSchemaWrapper - Matching Score Mechanism

#### The matching score system:
This is a temporary solution to how we target the correct ONNX overloads given that we only have manually annotated arguments (potentially inaccurate schema) and limited supports on AttributeProto.

1. Perfect match exam: If all arguments/kwargs are all matched, return the function without any warnings.
2. Best match exam: The system add the each correct matching input counts orderly, and subtract the symmetrical difference between their attributes to calculate the matching score. And select the one with the highest score in the end. If the selection is not a perfect match, a warning message is sent to SARIF.

#### Example of overloads

1. Different types: Caused by the difference between the ONNX spec and PyTorch.

The matching system finds the correct one.

```python
torch_op("aten::mul")
def aten_mul(self: TReal, other: TReal) -> TReal:
    ...

torch_op("aten::mul")
def aten_mul_bool(self: BOOL, other: BOOL) -> BOOL:
    ...
```

2. Optional dim: caused by unsupported op.OptionalHasElement (will support on opset version == 20). dim could be "None"

```python
torch_op("aten::argmax", trace_only=True)
def aten_argmax(
    self: TrealOrUInt8, dim: Optional[int] = None, keepdim: bool = False
) -> TrealOrUInt8:
    ...

torch_op("aten::argmax", private=True)
def _aten_argmax_dim(self: TrealOrUInt8, dim: int, keepdim: bool = False) -> TrealOrUInt8:
    ...
```

This case is impossible to differentiate, as they both might have dim in kwargs, so in this case, please make sure you turn the one with `dim: int` to private function.

3. Optional dtype: dtype could be "unprovided". The difference from 2 is that dtype would not be None.

```python
torch_op("aten::new_full")
def aten_new_full(self: TTensor, size: INT64, fill_value: TTensor) -> TTensor:
    ...

torch_op("aten::new_full")
def aten_new_full_dtype(self: TTensor, size: INT64, fill_value: TTensor, dtype: int) -> TTensor:
    ...
```

Depends on dtype is provided or not, matching system will dispatch the ATen op to the correct one.

4. `None` and `[]` and `NoneType` are considered failing the match. 

5. Two functions have the same score is recorded into SARIFs. 

### TODOs

1. Type promotion can be captured by dispatcher only if OpSchema can provide it. However, the implementation of "graph-level" pass vs "in-op"" promotion can be further discussed in microsoft/onnxscript#563.
5. torchlib should provide the "opset version" to OnnxRegistry.
7. How to expose OnnxRegistry with custom add/remove ops APIs nneds to be further discussed.

Co-authored-by: Justin Chu <justinchubymicrosoft.com>

[ghstack-poisoned]

Author: titaiwangms

.github/ci_commit_pins/torchbench.txt

@@ -1 +1 @@
-44ccad1d0647d56fe4e56c0b1933102be5dcc874
+f4acd1a7fcce986155c5e20beffa92b24ae0a3fa

.github/ci_commit_pins/vision.txt

@@ -1 +1 @@
-99ec261c72c097160e94653cce6f90f2d1209222
+d2f7486ccaef461913cdb51990ff353addf6f064

.github/scripts/label_utils.py

@@ -15,7 +15,7 @@
 
 BOT_AUTHORS = ["github-actions", "pytorchmergebot", "pytorch-bot"]
 
-LABEL_ERR_MSG_TITLE = "This PR needs a label"
+LABEL_ERR_MSG_TITLE = "This PR needs a `release notes:` label"
 LABEL_ERR_MSG = f"""# {LABEL_ERR_MSG_TITLE}
 If your changes are user facing and intended to be a part of release notes, please use a label starting with `release notes:`.
 
@@ -111,7 +111,9 @@ def has_required_labels(pr: "GitHubPR") -> bool:
 
 
 def is_label_err_comment(comment: GitHubComment) -> bool:
+    # comment.body_text returns text without markdown
+    no_format_title = LABEL_ERR_MSG_TITLE.replace("`", "")
     return (
-        comment.body_text.lstrip(" #").startswith(LABEL_ERR_MSG_TITLE)
+        comment.body_text.lstrip(" #").startswith(no_format_title)
         and comment.author_login in BOT_AUTHORS
     )

.github/scripts/test_check_labels.py

@@ -44,7 +44,7 @@ def mock_get_comments() -> List[GitHubComment]:
         ),
         # Case 2 - a label err comment
         GitHubComment(
-            body_text=" #" + LABEL_ERR_MSG_TITLE,
+            body_text=" #" + LABEL_ERR_MSG_TITLE.replace("`", ""),
             created_at="",
             author_login=BOT_AUTHORS[1],
             author_association="",

.github/scripts/trymerge.py

@@ -1262,9 +1262,9 @@ def find_matching_merge_rule(
                 reject_reason_score = num_matching_files
                 reject_reason = "\n".join(
                     (
-                        f"Not all files match rule `{rule_name}`."
-                        f"{num_matching_files} files matched, but there are still non-matching files:"
-                        f"{','.join(non_matching_files[:5])}{', ...' if len(non_matching_files) > 5 else ''}"
+                        f"Not all files match rule `{rule_name}`.",
+                        f"{num_matching_files} files matched, but there are still non-matching files:",
+                        f"{','.join(non_matching_files[:5])}{', ...' if len(non_matching_files) > 5 else ''}",
                     )
                 )
             continue

.github/scripts/update_commit_hashes.py

@@ -6,13 +6,13 @@
 
 import requests
 
-MERGEBOT_TOKEN = os.environ["MERGEBOT_TOKEN"]
+UPDATEBOT_TOKEN = os.environ["UPDATEBOT_TOKEN"]
 PYTORCHBOT_TOKEN = os.environ["PYTORCHBOT_TOKEN"]
 OWNER, REPO = "pytorch", "pytorch"
 
 
 def git_api(
-    url: str, params: Dict[str, str], type: str = "get", token: str = MERGEBOT_TOKEN
+    url: str, params: Dict[str, str], type: str = "get", token: str = UPDATEBOT_TOKEN
 ) -> Any:
     headers = {
         "Accept": "application/vnd.github.v3+json",

.github/workflows/_update-commit-hash.yml

@@ -22,7 +22,7 @@ on:
         required: false
         default: .github/ci_commit_pins
     secrets:
-      MERGEBOT_TOKEN:
+      UPDATEBOT_TOKEN:
         required: true
         description: Permissions for opening PR
       PYTORCHBOT_TOKEN:
@@ -41,7 +41,7 @@ jobs:
         with:
           fetch-depth: 1
           submodules: false
-          token: ${{ secrets.MERGEBOT_TOKEN }}
+          token: ${{ secrets.UPDATEBOT_TOKEN }}
 
       - name: Checkout
         shell: bash
@@ -54,11 +54,11 @@ jobs:
           REPO_NAME: ${{ inputs.repo-name }}
           BRANCH: ${{ inputs.branch }}
           PIN_FOLDER: ${{ inputs.pin-folder }}
-          MERGEBOT_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+          UPDATEBOT_TOKEN: ${{ secrets.UPDATEBOT_TOKEN }}
           PYTORCHBOT_TOKEN: ${{ secrets.PYTORCHBOT_TOKEN }}
         run: |
           # put this here instead of the script to prevent accidentally changing the config when running the script locally
-          git config --global user.name "PyTorch MergeBot"
-          git config --global user.email "pytorchmergebot@users.noreply.github.com"
+          git config --global user.name "PyTorch UpdateBot"
+          git config --global user.email "pytorchupdatebot@users.noreply.github.com"
 
           python .github/scripts/update_commit_hashes.py --repo-name "${REPO_NAME}" --branch "${BRANCH}" --pin-folder "${PIN_FOLDER}"

.github/workflows/check-labels.yml

@@ -1,8 +1,22 @@
 name: Check Labels
 
 on:
+  # We need pull_request_target to be able to post comments on PRs from forks.
+  # Only allow pull_request_target when merging to main, not some historical branch.
+  #
+  # Make sure to don't introduce explicit checking out and installing/running
+  # untrusted user code into this workflow!
+  pull_request_target:
+    types: [opened, synchronize, reopened, labeled, unlabeled]
+    branches: [main]
+    paths-ignore: [.github]
+
+  # To allow testing PRs that change workflows.
+  # May be triggered together with pull_request_target, it's OK.
   pull_request:
     types: [opened, synchronize, reopened, labeled, unlabeled]
+    paths: [.github]
+
   workflow_dispatch:
 
 concurrency:

.github/workflows/inductor.yml

@@ -73,7 +73,7 @@ jobs:
       docker-image-name: pytorch-linux-focal-py3.8-gcc7
       test-matrix: |
         { include: [
-          { config: "inductor_huggingface_cpu_accuracy", shard: 1, num_shards: 1, runner: "linux.4xlarge" },
+          { config: "inductor_huggingface_cpu_accuracy", shard: 1, num_shards: 1, runner: "linux.12xlarge" },
           { config: "inductor_timm_cpu_accuracy", shard: 1, num_shards: 2, runner: "linux.4xlarge" },
           { config: "inductor_timm_cpu_accuracy", shard: 2, num_shards: 2, runner: "linux.4xlarge" },
           { config: "inductor_torchbench_cpu_accuracy", shard: 1, num_shards: 1, runner: "linux.4xlarge" },

.github/workflows/lint-bc.yml

@@ -8,6 +8,8 @@ on:
       - reopened
       - labeled
       - unlabeled
+    branches-ignore:
+      - nightly
   workflow_dispatch:
 
 jobs:

.github/workflows/lint.yml

@@ -2,6 +2,8 @@ name: Lint
 
 on:
   pull_request:
+    branches-ignore:
+      - nightly
   push:
     branches:
       - main

.github/workflows/nightly.yml

@@ -43,5 +43,5 @@ jobs:
       repo-name: vision
       branch: main
     secrets:
-      MERGEBOT_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+      UPDATEBOT_TOKEN: ${{ secrets.UPDATEBOT_TOKEN }}
       PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN }}

.github/workflows/pull.yml

@@ -2,6 +2,8 @@ name: pull
 
 on:
   pull_request:
+    branches-ignore:
+      - nightly
   push:
     branches:
       - main

.github/workflows/revert.yml

@@ -8,6 +8,7 @@ jobs:
   do_revert:
     name: try_revert_pr_${{ github.event.client_payload.pr_num }}
     runs-on: linux.20_04.4x
+    environment: mergebot
     env:
         GH_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
     steps:

.github/workflows/trymerge.yml

@@ -8,6 +8,7 @@ jobs:
   do_merge:
     name: try_merge_pr_${{ github.event.client_payload.pr_num }}
     runs-on: linux.20_04.4x
+    environment: mergebot
     env:
         GH_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
     steps:

.github/workflows/tryrebase.yml

@@ -7,6 +7,7 @@ on:
 jobs:
   do_rebase:
     runs-on: ubuntu-20.04
+    environment: mergebot
     env:
         GH_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
     steps:

.github/workflows/update-viablestrict.yml

@@ -12,6 +12,7 @@ concurrency:
 jobs:
   do_update_viablestrict:
     runs-on: ubuntu-20.04
+    environment: mergebot
     steps:
       - name: Checkout repo
         uses: actions/checkout@v3

.github/workflows/upload-test-stats.yml

@@ -91,8 +91,6 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN}}
-          MERGEBOT_TOKEN: ${{ secrets.MERGEBOT_TOKEN}}
         run: |
           curl -H "Accept: application/vnd.github.v3+json" -H "Authorization: token $GITHUB_TOKEN" https://api.github.com/rate_limit
           curl -H "Accept: application/vnd.github.v3+json" -H "Authorization: token $PYTORCHBOT_TOKEN" https://api.github.com/rate_limit
-          curl -H "Accept: application/vnd.github.v3+json" -H "Authorization: token $MERGEBOT_TOKEN" https://api.github.com/rate_limit

.github/workflows/weekly.yml

@@ -15,7 +15,7 @@ jobs:
       repo-name: xla
       branch: master
     secrets:
-      MERGEBOT_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+      UPDATEBOT_TOKEN: ${{ secrets.UPDATEBOT_TOKEN }}
       PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN }}
 
   update-triton-commit-hash:
@@ -26,5 +26,5 @@ jobs:
       branch: main
       pin-folder: .ci/docker/ci_commit_pins
     secrets:
-      MERGEBOT_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
+      UPDATEBOT_TOKEN: ${{ secrets.UPDATEBOT_TOKEN }}
       PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN }}

.lintrunner.toml

@@ -1018,6 +1018,6 @@ init_command = [
     'python3',
     'tools/linter/adapters/pip_init.py',
     '--dry-run={{DRYRUN}}',
-    'ruff==0.0.265',
+    'ruff==0.0.269',
 ]
 is_formatter = true

BUILD.bazel

@@ -1696,6 +1696,36 @@ pybind_extension(
     ],
 )
 
+cc_library(
+    name = "functorch",
+    hdrs = glob([
+        "functorch/csrc/dim/*.h",
+    ]),
+    srcs = glob([
+        "functorch/csrc/dim/*.cpp",
+    ]),
+    deps = [
+        ":aten_nvrtc",
+        ":torch_python",
+        "@pybind11",
+    ],
+)
+
+pybind_extension(
+    name = "functorch/_C",
+    copts=[
+        "-DTORCH_EXTENSION_NAME=_C"
+    ],
+    srcs = [
+        "functorch/csrc/init_dim_only.cpp",
+    ],
+    deps = [
+        ":functorch",
+        ":torch_python",
+        ":aten_nvrtc",
+    ],
+)
+
 cc_binary(
     name = "torch/bin/torch_shm_manager",
     srcs = [
@@ -1724,7 +1754,7 @@ template_rule(
 rules.py_library(
     name = "pytorch_py",
     visibility = ["//visibility:public"],
-    srcs = glob(["torch/**/*.py"], exclude = ["torch/version.py"]) + [":torch/version.py"],
+    srcs = glob(["torch/**/*.py"], exclude = ["torch/version.py"]) + [":torch/version.py"] + glob(["functorch/**/*.py"]),
     deps = [
         rules.requirement("future"),
         rules.requirement("numpy"),
@@ -1737,6 +1767,7 @@ rules.py_library(
     ],
     data = [
         ":torch/_C.so",
+        ":functorch/_C.so",
         ":torch/bin/torch_shm_manager",
     ],
 )
@@ -1903,7 +1934,8 @@ cc_test(
 
 py_test(
     name = "test_bazel",
-    srcs = ["test/test_bazel.py"],
+    srcs = ["test/_test_bazel.py"],
+    main = "test/_test_bazel.py",
     deps = [":pytorch_py"],
 )
 

CONTRIBUTING.md

@@ -1228,7 +1228,7 @@ an active PR, CI jobs will be run automatically. Some of these may
 fail and you will need to find out why, by looking at the logs.
 
 Fairly often, a CI failure might be unrelated to your changes. You can
-confirm by going to our [HUD](hud.pytorch.org) and seeing if the CI job
+confirm by going to our [HUD](https://hud.pytorch.org) and seeing if the CI job
 is failing upstream already. In this case, you
 can usually ignore the failure. See [the following
 subsection](#which-commit-is-used-in-ci) for more details.

README.md

@@ -434,14 +434,14 @@ Three-pointers to get you started:
 
 ## Releases and Contributing
 
-PyTorch has a 90-day release cycle (major releases). Please let us know if you encounter a bug by [filing an issue](https://github.com/pytorch/pytorch/issues).
+Typically, PyTorch has three major releases a year. Please let us know if you encounter a bug by [filing an issue](https://github.com/pytorch/pytorch/issues).
 
 We appreciate all contributions. If you are planning to contribute back bug-fixes, please do so without any further discussion.
 
 If you plan to contribute new features, utility functions, or extensions to the core, please first open an issue and discuss the feature with us.
 Sending a PR without discussion might end up resulting in a rejected PR because we might be taking the core in a different direction than you might be aware of.
 
-To learn more about making a contribution to Pytorch, please see our [Contribution page](CONTRIBUTING.md).
+To learn more about making a contribution to Pytorch, please see our [Contribution page](CONTRIBUTING.md). For more information about PyTorch releases, see [Release page](RELEASE.md).
 
 ## The Team
 

aten/conda/meta.yaml

@@ -26,7 +26,7 @@ requirements:
 about:
   home: https://github.com/pytorch/pytorch
   license: BSD
-  summary: A TENsor library for C++14
+  summary: A TENsor library for C++17
 
 extra:
   recipe-maintainers:

aten/src/ATen/ATen.h

@@ -1,7 +1,7 @@
 #pragma once
 
-#if !defined(_MSC_VER) && __cplusplus < 201402L
-#error C++14 or later compatible compiler is required to use ATen.
+#if !defined(_MSC_VER) && __cplusplus < 201703L
+#error C++17 or later compatible compiler is required to use ATen.
 #endif
 
 #include <ATen/Context.h>

aten/src/ATen/LegacyVmapTransforms.h

@@ -40,6 +40,7 @@ using VmapPhysicalViewVec =
 // dimensions to get 8. Adjust this number as necessary
 constexpr int64_t kVmapStaticDimVecSize = 8;
 using VmapDimVector = SmallVector<int64_t, kVmapStaticDimVecSize>;
+using VmapSymDimVector = SmallVector<c10::SymInt, kVmapStaticDimVecSize>;
 
 // NOTE: [What is an VmapTransform?]
 // An *VmapTransform* converts logical views of tensors to physical views.

aten/src/ATen/cpu/tbb/CMakeLists.txt

@@ -121,9 +121,9 @@ endif()
 if(UNIX)
   add_definitions(-DUSE_PTHREAD)
 
-   check_cxx_compiler_flag("-std=c++14" SUPPORTS_STDCXX14)
-  if(SUPPORTS_STDCXX14)
-    set(CMAKE_CXX_FLAGS "-std=c++14 ${CMAKE_CXX_FLAGS}")
+   check_cxx_compiler_flag("-std=c++17" SUPPORTS_STDCXX17)
+  if(SUPPORTS_STDCXX17)
+    set(CMAKE_CXX_FLAGS "-std=c++17 ${CMAKE_CXX_FLAGS}")
   endif()
 
    check_cxx_compiler_flag("-mrtm -Werror" SUPPORTS_MRTM)

aten/src/ATen/cpu/vec/vec256/vsx/vec256_common_vsx.h

@@ -127,6 +127,30 @@ inline void convert(const int64_t* src, double* dst, int64_t n) {
     dst[i] = static_cast<double>(src[i]);
   }
 }
+//Generic implementation to fix compiler error
+//TO-DO : Add optimized version for ppc64
+inline std::tuple<Vectorized<float>, Vectorized<float>> convert_half_float(
+    const Vectorized<Half>& a) {
+  constexpr int64_t K = Vectorized<Half>::size();
+  __at_align__ float arr[K];
+  __at_align__ Half arr2[K];
+  a.store(arr2);
+  convert(arr2, arr, K);
+  return std::make_tuple(
+       Vectorized<float>::loadu(arr),
+       Vectorized<float>::loadu(arr + Vectorized<float>::size()));
+}
+
+inline Vectorized<Half> convert_float_half(
+    const Vectorized<float>& a, const Vectorized<float>& b) {
+  constexpr int64_t K = Vectorized<Half>::size();
+  __at_align__ float arr[K];
+  __at_align__ Half arr2[K];
+  a.store(arr);
+  b.store(arr + Vectorized<float>::size());
+  convert(arr, arr2, K);
+  return Vectorized<Half>::loadu(arr2);
+};
 
 template <>
 std::pair<Vectorized<double>, Vectorized<double>> inline interleave2<double>(