Help fix numpy detection in cross compiled layouts (#137084)

We had trouble at conda-forge getting numpy to get detected on aarch64 due to our splayed layout and cross compilation needs.

see:
* https://github.com/conda-forge/pytorch-cpu-feedstock/pull/256
* https://github.com/conda-forge/pytorch-cpu-feedstock/issues/266
* https://github.com/conda-forge/pytorch-cpu-feedstock/pull/267

This is my attempt at making an "upstreamable patch" that tries to follow your structure.

It could introduce a new environment variable `Python_NumPy_INCLUDE_DIR` if you want, but CMake doesn't use it as an environment variable, so I feel like that would be weird.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/137084
Approved by: https://github.com/atalman

tools/setup_helpers/cmake.py

@@ -338,6 +338,7 @@ class CMake:
         # future, as CMake can detect many of these libraries pretty comfortably. We have them here for now before CMake
         # integration is completed. They appear here not in the CMake.defines call below because they start with either
         # "BUILD_" or "USE_" and must be overwritten here.
+        use_numpy = not check_negative_env_flag("USE_NUMPY")
         build_options.update(
             {
                 # Note: Do not add new build options to this dict if it is directly read from environment variable -- you
@@ -347,7 +348,7 @@ class CMake:
                 "BUILD_TEST": build_test,
                 # Most library detection should go to CMake script, except this one, which Python can do a much better job
                 # due to NumPy's inherent Pythonic nature.
-                "USE_NUMPY": not check_negative_env_flag("USE_NUMPY"),
+                "USE_NUMPY": use_numpy,
             }
         )
 
@@ -373,6 +374,20 @@ class CMake:
             sys.exit(1)
         build_options.update(cmake__options)
 
+        if use_numpy:
+            try:
+                # This helps CMake find the correct include directory for NumPy
+                # This is especially useful in cross compiled environments
+                import numpy
+
+                Python_NumPy_INCLUDE_DIR = numpy.get_include()
+                build_options.update(
+                    dict(Python_NumPy_INCLUDE_DIR=Python_NumPy_INCLUDE_DIR)
+                )
+            except ImportError:
+                # use_numpy is just a hint.... so we can fail silently here
+                pass
+
         CMake.defines(
             args,
             Python_EXECUTABLE=sys.executable,