Add a stable TORCH_LIBRARY to C shim (#148124)

This PR adds two main parts:
- shim.h stable C APIs into torch::Library APIs
- a higher level API in torch/csrc/stable/library.h that calls into this shim.h + otherwise is self contained

Goal: custom kernel writers should be able to call the apis in the directories above in order to register their library in a way that allows their custom extension to run with a different libtorch version than it was built with.

Subplots resolved:

- Do we want a whole separate StableLibrary or do we want to freeze torch::Library and add `m.stable_impl(cstring, void (*fn)(void **, int64_t, int64_t)` into it
    - Yes, we want a separate StableLibrary. We cannot freeze Library and it is NOT header only.
- Should I use unint64_t as the common denominator instead of void* to support 32bit architectures better?
    -  Yes, and done
- Should I add a stable `def` and `fragment` when those can be done in python?
    - I think we do want these --- and now they're done
- Where should library_stable_impl.cpp live? -- no longer relevant
- I need some solid test cases to make sure everything's going ok. I've intentionally thrown in a bunch of random dtypes into the signature, but I still haven't tested returning multiple things, returning nothing, complex dtypes, etc.
    - Have since tested all the torch library endpoints. the others can be tested in a followup to separate components that need to be in shim.h vs can be added later

Pull Request resolved: https://github.com/pytorch/pytorch/pull/148124
Approved by: https://github.com/albanD, https://github.com/zou3519, https://github.com/atalman

test/run_test.py

@@ -931,10 +931,10 @@ def install_cpp_extensions(cpp_extensions_test_dir, env=os.environ):
 
 
 @contextlib.contextmanager
-def extend_python_path(install_directory):
+def extend_python_path(install_directories):
     python_path = os.environ.get("PYTHONPATH", "")
     try:
-        os.environ["PYTHONPATH"] = os.pathsep.join([install_directory, python_path])
+        os.environ["PYTHONPATH"] = os.pathsep.join(install_directories + [python_path])
         yield
     finally:
         os.environ["PYTHONPATH"] = python_path
@@ -1074,9 +1074,12 @@ def _test_cpp_extensions_aot(test_directory, options, use_ninja):
     if return_code != 0:
         return return_code
     if sys.platform != "win32":
-        exts_to_build = [(install_cmd, "no_python_abi_suffix_test")]
+        exts_to_build = [
+            (install_cmd, "no_python_abi_suffix_test"),
+        ]
         if TEST_CUDA:
             exts_to_build.append((wheel_cmd, "python_agnostic_extension"))
+            exts_to_build.append((install_cmd, "libtorch_agnostic_extension"))
         for cmd, extension_dir in exts_to_build:
             return_code = shell(
                 cmd,
@@ -1094,17 +1097,24 @@ def _test_cpp_extensions_aot(test_directory, options, use_ninja):
         test_directory + "/test_cpp_extensions_aot.py",
         test_directory + "/" + test_module + ".py",
     )
+
     try:
         cpp_extensions = os.path.join(test_directory, "cpp_extensions")
-        install_directory = ""
+        install_directories = []
         # install directory is the one that is named site-packages
         for root, directories, _ in os.walk(os.path.join(cpp_extensions, "install")):
             for directory in directories:
                 if "-packages" in directory:
-                    install_directory = os.path.join(root, directory)
+                    install_directories.append(os.path.join(root, directory))
 
-        assert install_directory, "install_directory must not be empty"
-        with extend_python_path(install_directory):
+        for root, directories, _ in os.walk(
+            os.path.join(cpp_extensions, "libtorch_agnostic_extension", "install")
+        ):
+            for directory in directories:
+                if "-packages" in directory:
+                    install_directories.append(os.path.join(root, directory))
+
+        with extend_python_path(install_directories):
             return run_test(ShardedTest(test_module, 1, 1), test_directory, options)
     finally:
         if os.path.exists(test_directory + "/" + test_module + ".py"):
@@ -1136,7 +1146,7 @@ def _test_autoload(test_directory, options, enable=True):
 
     try:
         os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = str(int(enable))
-        with extend_python_path(install_directory):
+        with extend_python_path([install_directory]):
             cmd = [sys.executable, "test_autoload.py"]
             return_code = shell(cmd, cwd=test_directory, env=os.environ)
             return return_code
@@ -1152,7 +1162,7 @@ def run_test_with_openreg(test_module, test_directory, options):
     if return_code != 0:
         return return_code
 
-    with extend_python_path(install_dir):
+    with extend_python_path([install_dir]):
         return run_test(test_module, test_directory, options)