From cb5e1fd234c9f7c2c5b36024d136a9dd18b0fbe8 Mon Sep 17 00:00:00 2001
From: Guillaume Fraux <guillaume.fraux@epfl.ch>
Date: Fri, 29 Nov 2024 16:11:01 +0100
Subject: [PATCH] Make sure shm.dll can find mkl.dll in torch 2.3+

---
 .github/workflows/build-wheels.yml | 20 +++++++++----------
 python/featomic_torch/setup.py     | 32 ++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml
index fe4f9a227..d4c4926cd 100644
--- a/.github/workflows/build-wheels.yml
+++ b/.github/workflows/build-wheels.yml
@@ -95,7 +95,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: featomic-wheel-${{ matrix.os }}-${{ matrix.cibw-arch }}
+          name: wheel-${{ matrix.os }}-${{ matrix.cibw-arch }}
           path: ./wheelhouse/*.whl
 
   build-torch-wheels:
@@ -214,13 +214,13 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: featomic-torch-single-version-wheel-${{ matrix.torch-version }}-${{ matrix.os }}-${{ matrix.arch }}
+          name: torch-single-version-wheel-${{ matrix.torch-version }}-${{ matrix.os }}-${{ matrix.arch }}
           path: ./wheelhouse/*.whl
 
   merge-torch-wheels:
     needs: build-torch-wheels
     runs-on: ubuntu-22.04
-    name: merge metatensor-torch ${{ matrix.name }}
+    name: merge featomic-torch ${{ matrix.name }}
     strategy:
       matrix:
         include:
@@ -287,10 +287,10 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: featomic-torch-wheel-${{ matrix.os }}-${{ matrix.arch }}
+          name: torch-wheel-${{ matrix.os }}-${{ matrix.arch }}
           path: ./wheelhouse/*.whl
 
-  build-others:
+  build-sdist:
     runs-on: ubuntu-22.04
     name: sdist
     steps:
@@ -319,7 +319,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: others
+          name: sdist
           path: |
             dist/*.tar.gz
             dist/cxx/*.tar.gz
@@ -336,7 +336,7 @@ jobs:
 
   merge-and-release:
     name: Merge and release wheels/sdists
-    needs: [build-wheels, merge-torch-wheels, build-others]
+    needs: [build-wheels, merge-torch-wheels, build-sdist]
     runs-on: ubuntu-22.04
     permissions:
       contents: write
@@ -345,21 +345,21 @@ jobs:
         uses: actions/download-artifact@v4
         with:
           path: wheels
-          pattern: featomic-wheel-*
+          pattern: wheel-*
           merge-multiple: true
 
       - name: Download featomic-torch wheels
         uses: actions/download-artifact@v4
         with:
           path: wheels
-          pattern: featomic-torch-wheel-*
+          pattern: torch-wheel-*
           merge-multiple: true
 
       - name: Download sdists
         uses: actions/download-artifact@v4
         with:
           path: wheels
-          name: others
+          name: sdist
 
       - name: Re-upload a single wheels artifact
         uses: actions/upload-artifact@v4
diff --git a/python/featomic_torch/setup.py b/python/featomic_torch/setup.py
index fbd14c4bd..3d80532db 100644
--- a/python/featomic_torch/setup.py
+++ b/python/featomic_torch/setup.py
@@ -268,6 +268,38 @@ def create_version_number(version):
 
 
 if __name__ == "__main__":
+    if sys.platform == "win32":
+        # On Windows, starting with PyTorch 2.3, the file shm.dll in torch has a
+        # dependency on mkl DLLs. When building the code using pip build isolation, pip
+        # installs the mkl package in a place where the os is not trying to load
+        #
+        # This is a very similar fix to https://github.com/pytorch/pytorch/pull/126095,
+        # except only applying when importing torch from a build-isolation virtual
+        # environment created by pip (`python -m build` does not seems to suffer from
+        # this).
+        import wheel
+
+        pip_virtualenv = os.path.realpath(
+            os.path.join(
+                os.path.dirname(wheel.__file__),
+                "..",
+                "..",
+                "..",
+                "..",
+            )
+        )
+        mkl_dll_dir = os.path.join(
+            pip_virtualenv,
+            "normal",
+            "Library",
+            "bin",
+        )
+
+        if os.path.exists(mkl_dll_dir):
+            os.add_dll_directory(mkl_dll_dir)
+
+        # End of Windows/MKL/PIP hack
+
     if not os.path.exists(FEATOMIC_TORCH_SRC):
         # we are building from a sdist, which should include featomic-torch
         # sources as a tarball