Enable pymomentum on Windows

.github/workflows/ci_ubuntu.yml

@@ -100,6 +100,6 @@ jobs:
           environments: ${{ matrix.pixi_env }}
           cache: true
 
-      - name: Build PyMomentum
+      - name: Build and test PyMomentum
         run: |
           pixi run -e ${{ matrix.pixi_env }} test_py

.github/workflows/ci_windows.yml

@@ -13,7 +13,7 @@ on:
       - "**/website/**"
 
 jobs:
-  build:
+  momentum:
     name: cpp-${{ matrix.mode == '' && 'opt' || 'dev' }}-win
     runs-on: windows-latest
     strategy:
@@ -43,3 +43,19 @@ jobs:
             --build momentum/examples/hello_world/build `
             --config Release `
             --parallel
+
+  pymomentum:
+    name: py-ubuntu
+    runs-on: windows-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up pixi
+        uses: prefix-dev/[email protected]
+        with:
+          cache: true
+
+      - name: Build and test PyMomentum
+        run: |
+          pixi run test_py

pixi.toml

@@ -39,6 +39,7 @@ ms-gsl = ">=4.1.0,<5"
 nlohmann_json = ">=3.11.3,<4"
 openfbx = ">=0.9,<0.10"
 openssl = ">=3.4.0,<4"
+pytorch = ">=2.5.1,<3"
 re2 = ">=2024.7.2,<2025"
 spdlog = ">=1.15.0,<2"
 
@@ -104,6 +105,18 @@ install = { cmd = "cmake --build build/$PIXI_ENVIRONMENT_NAME/cpp/release -j --t
     "build",
 ] }
 
+test_py = { cmd = """
+    pytest \
+        pymomentum/test/test_closest_points.py \
+        pymomentum/test/test_fbx.py \
+        pymomentum/test/test_parameter_transform.py \
+        pymomentum/test/test_quaternion.py \
+        pymomentum/test/test_skel_state.py \
+        pymomentum/test/test_skeleton.py
+    """, env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
+    "build_py",
+] }
+
 #===========
 # linux-64
 #===========
@@ -112,7 +125,6 @@ install = { cmd = "cmake --build build/$PIXI_ENVIRONMENT_NAME/cpp/release -j --t
 nvtx-c = ">=3.1.0" # TODO: Add to pytorch as run dep
 
 [target.linux-64.dependencies]
-pytorch = ">=2.4.0"
 sysroot_linux-64 = ">=2.28"
 
 [target.linux-64.tasks]
@@ -173,17 +185,6 @@ build_py = { cmd = "pip install . -vv", env = { FBXSDK_PATH = ".deps/fbxsdk", CM
 """, MOMENTUM_ENABLE_SIMD = "ON" }, depends-on = [
     "install_deps",
 ] }
-test_py = { cmd = """
-    pytest \
-        pymomentum/test/test_closest_points.py \
-        pymomentum/test/test_fbx.py \
-        pymomentum/test/test_parameter_transform.py \
-        pymomentum/test/test_quaternion.py \
-        pymomentum/test/test_skel_state.py \
-        pymomentum/test/test_skeleton.py
-    """, env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
-    "build_py",
-] }
 
 #============
 # osx-64
@@ -192,24 +193,12 @@ test_py = { cmd = """
 [target.osx.build-dependencies]
 
 [target.osx.dependencies]
-pytorch = ">=2.4.0"
 
 [target.osx.tasks]
 build_py = { cmd = "pip install . -vv", env = { CMAKE_ARGS = """
     -DMOMENTUM_BUILD_IO_FBX=$MOMENTUM_BUILD_IO_FBX \
     -DMOMENTUM_ENABLE_SIMD=$MOMENTUM_ENABLE_SIMD
 """, MOMENTUM_BUILD_IO_FBX = "OFF", MOMENTUM_ENABLE_SIMD = "ON" } }
-test_py = { cmd = """
-    pytest \
-        pymomentum/test/test_closest_points.py \
-        pymomentum/test/test_fbx.py \
-        pymomentum/test/test_parameter_transform.py \
-        pymomentum/test/test_quaternion.py \
-        pymomentum/test/test_skel_state.py \
-        pymomentum/test/test_skeleton.py
-    """, env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
-    "build_py",
-] }
 
 #============
 # osx-arm64
@@ -218,30 +207,19 @@ test_py = { cmd = """
 [target.osx-arm64.build-dependencies]
 
 [target.osx-arm64.dependencies]
-pytorch = ">=2.4.0"
 
 [target.osx-arm64.tasks]
 build_py = { cmd = "pip install . -vv", env = { CMAKE_ARGS = """
     -DMOMENTUM_BUILD_IO_FBX=$MOMENTUM_BUILD_IO_FBX \
     -DMOMENTUM_ENABLE_SIMD=$MOMENTUM_ENABLE_SIMD
 """, MOMENTUM_BUILD_IO_FBX = "OFF", MOMENTUM_ENABLE_SIMD = "ON" } }
-test_py = { cmd = """
-    pytest \
-        pymomentum/test/test_closest_points.py \
-        pymomentum/test/test_fbx.py \
-        pymomentum/test/test_parameter_transform.py \
-        pymomentum/test/test_quaternion.py \
-        pymomentum/test/test_skel_state.py \
-        pymomentum/test/test_skeleton.py
-    """, env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
-    "build_py",
-] }
 
 #=========
 # win-64
 #=========
 
 [target.win-64.dependencies]
+cuda-nvtx-dev = ">=12.6.77,<13"
 
 [target.win-64.tasks]
 config = { cmd = """
@@ -300,12 +278,17 @@ install = { cmd = "cmake --build build/$PIXI_ENVIRONMENT_NAME/cpp -j --target in
     "build",
 ] }
 
+build_py = { cmd = "pip install . -vv", env = { CMAKE_ARGS = """
+    -DMOMENTUM_BUILD_IO_FBX=$MOMENTUM_BUILD_IO_FBX \
+    -DMOMENTUM_ENABLE_SIMD=$MOMENTUM_ENABLE_SIMD
+""", MOMENTUM_BUILD_IO_FBX = "OFF", MOMENTUM_ENABLE_SIMD = "ON" } }
+
 #==============
 # Feature: CPU
 #==============
 
 [feature.cpu]
-platforms = ["win-64", "linux-64", "osx-64", "osx-arm64"]
+platforms = ["linux-64", "osx-64", "osx-arm64", "win-64"]
 
 #==============
 # Feature: GPU
@@ -315,19 +298,19 @@ platforms = ["win-64", "linux-64", "osx-64", "osx-arm64"]
 channels = ["conda-forge"]
 platforms = ["linux-64"]
 system-requirements = { cuda = "12.0" }
-dependencies = { cuda-toolkit = ">=12.6.3,<13", pytorch = { version = ">=2.5.1,<3", build = "cuda126_py312*", channel = "conda-forge" } }
+dependencies = { cuda-toolkit = "12.*", python = "3.12.*", pytorch-gpu = { version = ">=2.5.1,<3", build = "cuda126_*" } , cuda-nvtx-dev = ">=12.6.77,<13" }
 
 [feature.py311-cuda126]
 channels = ["conda-forge"]
 platforms = ["linux-64"]
 system-requirements = { cuda = "12.0" }
-dependencies = { cuda-toolkit = ">=12.6.3,<13", pytorch = { version = ">=2.5.1,<3", build = "cuda126_py311*", channel = "conda-forge" } }
+dependencies = { cuda-toolkit = "12.*", python = "3.11.*", pytorch-gpu = { version = ">=2.5.1,<3", build = "cuda126_*" } , cuda-nvtx-dev = ">=12.6.77,<13" }
 
 [feature.py310-cuda126]
 channels = ["conda-forge"]
 platforms = ["linux-64"]
 system-requirements = { cuda = "12.0" }
-dependencies = { cuda-toolkit = ">=12.6.3,<13", pytorch = { version = ">=2.5.1,<3", build = "cuda126_py310*", channel = "conda-forge" } }
+dependencies = { cuda-toolkit = "12.*", python = "3.10.*", pytorch-gpu = { version = ">=2.5.1,<3", build = "cuda126_*" } , cuda-nvtx-dev = ">=12.6.77,<13" }
 
 #==============
 # Environments

pymomentum/CMakeLists.txt

@@ -7,24 +7,31 @@
 # Find dependencies
 #===============================================================================
 
+if(NOT DEFINED ENV{CONDA_PREFIX})
+  message(FATAL_ERROR
+    "CONDA_PREFIX is not set. pymomentum currently only supports building"
+    "with Conda/Pixi. Please ensure you are using a Conda or Pixi environment."
+  )
+endif()
+
 set(ENV{NVTOOLSEXT_PATH} "$ENV{CONDA_PREFIX}/include")
 
-find_package(ATen CONFIG REQUIRED
-  HINTS
-    $ENV{CONDA_PREFIX}/lib/python${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}/site-packages/torch/
-)
+if(WIN32)
+  set(libtorch_base_path $ENV{CONDA_PREFIX}/Lib/site-packages/torch)
+else()  
+  set(libtorch_base_path $ENV{CONDA_PREFIX}/lib/python${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}/site-packages/torch)
+endif()
 
-find_package(Torch CONFIG REQUIRED
-  HINTS
-    $ENV{CONDA_PREFIX}/lib/python${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}/site-packages/torch/
-)
+if(NOT EXISTS "${libtorch_base_path}")
+  message(FATAL_ERROR
+    "PyTorch not found in the expected location: ${libtorch_base_path}"
+    "Please ensure PyTorch is installed in your Conda/Pixi environment."
+  )
+endif()
 
-find_library(torch_python
-  NAMES torch_python
-  HINTS
-    $ENV{CONDA_PREFIX}/lib/python${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}/site-packages/torch/lib/
-  REQUIRED
-)
+find_package(ATen CONFIG REQUIRED HINTS ${libtorch_base_path})
+find_package(Torch CONFIG REQUIRED HINTS ${libtorch_base_path})
+find_library(torch_python NAMES torch_python HINTS ${libtorch_base_path}/lib/ REQUIRED)
 
 #===============================================================================
 # Find build dependencies
@@ -82,8 +89,6 @@ mt_library(
   PYMOMENTUM_SOURCES_VARS tensor_momentum_sources
   PUBLIC_INCLUDE_DIRECTORIES
     ${ATEN_INCLUDE_DIR}
-  PRIVATE_INCLUDE_DIRECTORIES
-    ${TORCH_LIBRARIES}
   PUBLIC_LINK_LIBRARIES
     momentum
     ${ATEN_LIBRARIES}
@@ -108,7 +113,6 @@ mt_python_binding(
   PYMOMENTUM_SOURCES_VARS geometry_sources
   INCLUDE_DIRECTORIES
     ${ATEN_INCLUDE_DIR}
-    ${TORCH_LIBRARIES}
   LINK_LIBRARIES
     character
     character_test_helpers
@@ -134,7 +138,6 @@ mt_python_binding(
   PYMOMENTUM_SOURCES_VARS quaternion_sources
   INCLUDE_DIRECTORIES
     ${ATEN_INCLUDE_DIR}
-    ${TORCH_LIBRARIES}
   LINK_LIBRARIES
     tensor_momentum
     ${ATEN_LIBRARIES}
@@ -150,7 +153,6 @@ mt_python_binding(
   PYMOMENTUM_SOURCES_VARS skel_state_sources
   INCLUDE_DIRECTORIES
     ${ATEN_INCLUDE_DIR}
-    ${TORCH_LIBRARIES}
   LINK_LIBRARIES
     tensor_momentum
     ${ATEN_LIBRARIES}

pymomentum/tensor_momentum/tensor_skeleton_state.cpp

@@ -112,8 +112,7 @@ void computeSkelStateBackward(
       for (int d = 0; d < momentum::kParametersPerJoint; ++d) {
         Eigen::VectorX<JetType> jointParams_cur =
             jointParameters
-                .template segment(
-                    kParametersPerJoint * curJoint, kParametersPerJoint)
+                .segment(kParametersPerJoint * curJoint, kParametersPerJoint)
                 .template cast<JetType>();
         jointParams_cur(d).v[0] = 1;
 
@@ -189,8 +188,7 @@ void computeLocalSkelStateBackward(
     for (int d = 0; d < momentum::kParametersPerJoint; ++d) {
       Eigen::VectorX<JetType> jointParams_cur =
           jointParameters
-              .template segment(
-                  kParametersPerJoint * iJoint, kParametersPerJoint)
+              .segment(kParametersPerJoint * iJoint, kParametersPerJoint)
               .template cast<JetType>();
       jointParams_cur(d).v[0] = 1;
 

pymomentum/tensor_utility/tensor_utility.h

@@ -10,6 +10,12 @@
 #include <ATen/ATen.h>
 #include <torch/torch.h>
 #include <Eigen/Core>
+
+#if defined(_MSC_VER)
+#include <BaseTsd.h>
+using ssize_t = SSIZE_T;
+#endif
+
 #include <optional>
 #include <unordered_map>