Build pymomentum on Windows (#238)

Summary:
Now that the [PyTorch conda package is available on Windows](conda-forge/pytorch-cpu-feedstock#231), we can enable building PyMomentum on this platform.

## Checklist:

- [x] Adheres to the [style guidelines](https://facebookincubator.github.io/momentum/docs/developer_guide/style_guide)
- [x] Codebase formatted by running `pixi run lint`

Pull Request resolved: #238

Test Plan: CI on Windows

Reviewed By: juliencbmeta

Differential Revision: D71858372

Pulled By: jeongseok-meta

fbshipit-source-id: a05d8d88c7fcd81472ff787349ef9deef6f3deb9

Author: jeongseok-meta

.github/workflows/ci_macos.yml

@@ -20,7 +20,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [macos-latest, macos-latest-large]
-        mode: ["", "_dev"]
+        mode: [""]
     steps:
       - name: Checkout
         uses: actions/checkout@v4

.github/workflows/ci_ubuntu.yml

@@ -100,7 +100,7 @@ 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,48 @@ jobs:
             --build momentum/examples/hello_world/build `
             --config Release `
             --parallel
+
+  pymomentum:
+    name: py-${{ matrix.pixi_env }}-win
+    runs-on: windows-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - pixi_env: "cpu"
+          - pixi_env: "gpu"
+            cuda_version: "12.8.0"
+    env:
+      FULL_CUDA_VERSION: ${{ matrix.cuda_version }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Install CUDA Toolkit
+        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
+        uses: Jimver/cuda-toolkit@v0.2.22
+        id: cuda-toolkit
+        with:
+          # Available versions: https://github.com/Jimver/cuda-toolkit/blob/v0.2.21/src/links/linux-links.ts
+          cuda: ${{ matrix.cuda_version }}
+
+      - name: Check CUDA Version
+        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
+        run: |
+          nvcc --version
+
+      - name: Set Conda environment variables
+        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
+        shell: pwsh
+        run: |
+          $majorVersion = $Env:FULL_CUDA_VERSION.Split('.')[0]
+          "CONDA_OVERRIDE_CUDA=$majorVersion" | Out-File -FilePath $Env:GITHUB_ENV -Encoding utf8 -Append
+
+      - name: Set up pixi
+        uses: prefix-dev/setup-pixi@v0.8.1
+        with:
+          cache: true
+
+      - name: Build and test PyMomentum
+        run: |
+          pixi run -e ${{ matrix.pixi_env }} test_py

CMakeLists.txt

@@ -111,7 +111,7 @@ if(MOMENTUM_USE_SYSTEM_RERUN_CPP_SDK)
 else()
   include(FetchContent)
   FetchContent_Declare(rerun_sdk
-    URL https://github.com/rerun-io/rerun/releases/download/0.19.1/rerun_cpp_sdk.zip
+    URL https://github.com/rerun-io/rerun/releases/download/0.21.0/rerun_cpp_sdk.zip
   )
   FetchContent_MakeAvailable(rerun_sdk)
 endif()
@@ -496,7 +496,6 @@ mt_library(
   PUBLIC_LINK_LIBRARIES
     Eigen3::Eigen
     rerun_sdk
-  NO_INSTALL
 )
 
 mt_library(
@@ -511,7 +510,6 @@ mt_library(
   PRIVATE_LINK_LIBRARIES
     rerun_eigen_adapters
     axel
-  NO_INSTALL
 )
 
 #===============================================================================

cmake/mt_defs.cmake

@@ -351,7 +351,7 @@ function(mt_setup_gtest)
   )
 
   if(NOT _ARG_GIT_TAG)
-    set(_ARG_GIT_TAG v1.15.2)
+    set(_ARG_GIT_TAG v1.16.0)
   endif()
 
   if(MOMENTUM_USE_SYSTEM_GOOGLETEST)

momentum/io/urdf/urdf_io.cpp

@@ -82,9 +82,9 @@ bool loadUrdfSkeletonRecursive(
   joint.name = urdfLink->name; // Use link name or joint name?
   joint.parent = parentJointId;
 
-  const size_t jointId = data.skeleton.joints.size();
-  const size_t jointParamsBaseIndex = jointId * kParametersPerJoint;
-  const size_t modelParamsBaseIndex = data.totalDoFs;
+  const Eigen::Index jointId = data.skeleton.joints.size();
+  const Eigen::Index jointParamsBaseIndex = jointId * kParametersPerJoint;
+  const Eigen::Index modelParamsBaseIndex = data.totalDoFs;
 
   //---------------------------
   // Parse Parameter transform

pixi.lock

@@ -10,38 +10,39 @@ homepage = "https://facebookincubator.github.io/momentum/"
 repository = "https://github.com/facebookincubator/momentum"
 
 [build-dependencies]
-boost = ">=1.85.0,<2"
-c-compiler = ">=1.8.0,<2"
-clang-format = ">=18.1.8,<19"
-cmake = ">=3.31.2,<4"
-cxx-compiler = ">=1.8.0,<2"
-gtest = ">=1.15.2,<2"
+boost = ">=1.84.0,<2"
+c-compiler = ">=1.9.0,<2"
+clang-format = ">=20.1.1,<21"
+cmake = ">=3.31.6,<4"
+cxx-compiler = ">=1.9.0,<2"
+gtest = ">=1.16.0,<2"
 ninja = ">=1.12.1,<2"
 rerun-sdk = ">=0.21.0,<0.22"
-pip = ">=24.3.1,<25"
+pip = ">=25.0.1,<26"
 pybind11 = ">=2.13.6,<3"
-pytest = ">=8.3.4,<9"
-scipy = ">=1.15.0,<2"
-setuptools = ">=75.6.0,<76"
+pytest = ">=8.3.5,<9"
+scipy = ">=1.15.2,<2"
+setuptools = ">=75.8.2,<76"
 sphinx = ">=8.1.3,<9"
 sphinx-rtd-theme = ">=3.0.1,<4"
 tracy-profiler-gui = ">=0.11.1,<0.12"
 
 [dependencies]
-blas = ">=2.129,<3"
+blas = ">=2.131,<3"
 ceres-solver = ">=2.2.0,<3"
-cli11 = ">=2.4.2,<3"
+cli11 = ">=2.5.0,<3"
 dispenso = ">=1.4.0,<2"
 eigen = ">=3.4.0,<4"
-ezc3d = ">=1.5.17,<2"
+ezc3d = ">=1.5.18,<2"
 drjit-cpp = ">=1.0.5,<2"
-fmt = ">=11.0.2,<12"
+fmt = ">=11.1.4,<12"
 fx-gltf = ">=2.0.0,<3"
 librerun-sdk = ">=0.21.0,<0.22"
-ms-gsl = ">=4.1.0,<5"
+ms-gsl = ">=4.2.0,<5"
 nlohmann_json = ">=3.11.3,<4"
 openfbx = ">=0.9,<0.10"
 openssl = ">=3.4.1,<4"
+pytorch = ">=2.6.0,<3"
 re2 = ">=2024.7.2,<2025"
 spdlog = ">=1.15.1,<2"
 tracy-profiler-client = ">=0.11.1,<0.12"
@@ -132,6 +133,15 @@ install = { cmd = "cmake --build build/$PIXI_ENVIRONMENT_NAME/cpp/release -j --t
     "build",
 ] }
 tracy = { cmd = "tracy-profiler" }
+test_py = { cmd = "pytest pymomentum/test/*.py", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
+    "build_py",
+] }
+doc_py = { cmd = "sphinx-build -a -E -b html pymomentum/doc build/python_api_doc", depends-on = [
+    "build_py",
+] }
+open_doc_py = { cmd = "open build/python_api_doc/index.html", depends-on = [
+    "doc_py",
+] }
 
 #===========
 # linux-64
@@ -141,7 +151,6 @@ tracy = { cmd = "tracy-profiler" }
 nvtx-c = ">=3.1.0" # TODO: Add to pytorch as run dep
 
 [target.linux-64.dependencies]
-pytorch = ">=2.6.0,<3"
 sysroot_linux-64 = ">=2.28"
 
 [target.linux-64.tasks]
@@ -204,16 +213,18 @@ config_dev = { cmd = """
 ] }
 build_py = { cmd = "pip install . -vv", env = { FBXSDK_PATH = ".deps/fbxsdk", CMAKE_ARGS = """
     -DMOMENTUM_BUILD_IO_FBX=ON \
-    -DMOMENTUM_ENABLE_SIMD=$MOMENTUM_ENABLE_SIMD
+    -DMOMENTUM_ENABLE_SIMD=$MOMENTUM_ENABLE_SIMD \
+    -DMOMENTUM_USE_SYSTEM_GOOGLETEST=ON \
+    -DMOMENTUM_USE_SYSTEM_PYBIND11=OFF \
+    -DMOMENTUM_USE_SYSTEM_RERUN_CPP_SDK=ON
 """, MOMENTUM_ENABLE_SIMD = "ON" }, depends-on = [
     "install_deps",
 ] }
 test_py = { cmd = "pytest pymomentum/test/*.py", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
     "build_py",
 ] }
-install_py = { cmd = "pip install -e ." }
 doc_py = { cmd = "sphinx-build -a -E -b html pymomentum/doc build/python_api_doc", depends-on = [
-    "install_py",
+    "build_py",
 ] }
 open_doc_py = { cmd = "open build/python_api_doc/index.html", depends-on = [
     "doc_py",
@@ -226,20 +237,21 @@ open_doc_py = { cmd = "open build/python_api_doc/index.html", depends-on = [
 [target.osx.build-dependencies]
 
 [target.osx.dependencies]
-pytorch = ">=2.6.0,<3"
 
 [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
+    -DMOMENTUM_ENABLE_SIMD=$MOMENTUM_ENABLE_SIMD \
+    -DMOMENTUM_USE_SYSTEM_GOOGLETEST=ON \
+    -DMOMENTUM_USE_SYSTEM_PYBIND11=ON \
+    -DMOMENTUM_USE_SYSTEM_RERUN_CPP_SDK=ON
 """, MOMENTUM_BUILD_IO_FBX = "OFF", MOMENTUM_ENABLE_SIMD = "ON" } }
 # TODO: Remove -k option, once this tests are disabled programmatically if momentum is not built with FBX support
 test_py = { cmd = "pytest pymomentum/test/*.py -k 'not (TestFBXIO and (test_save_motions_with_joint_params or test_save_motions_with_model_params))'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
     "build_py",
 ] }
-install_py = { cmd = "pip install -e ." }
 doc_py = { cmd = "sphinx-build -a -E -b html pymomentum/doc build/python_api_doc", depends-on = [
-    "install_py",
+    "build_py",
 ] }
 open_doc_py = { cmd = "open build/python_api_doc/index.html", depends-on = [
     "doc_py",
@@ -252,20 +264,21 @@ open_doc_py = { cmd = "open build/python_api_doc/index.html", depends-on = [
 [target.osx-arm64.build-dependencies]
 
 [target.osx-arm64.dependencies]
-pytorch = ">=2.6.0,<3"
 
 [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
+    -DMOMENTUM_ENABLE_SIMD=$MOMENTUM_ENABLE_SIMD \
+    -DMOMENTUM_USE_SYSTEM_GOOGLETEST=ON \
+    -DMOMENTUM_USE_SYSTEM_PYBIND11=ON \
+    -DMOMENTUM_USE_SYSTEM_RERUN_CPP_SDK=ON
 """, MOMENTUM_BUILD_IO_FBX = "OFF", MOMENTUM_ENABLE_SIMD = "ON" } }
 # TODO: Remove -k option, once this tests are disabled programmatically if momentum is not built with FBX support
 test_py = { cmd = "pytest pymomentum/test/*.py -k 'not (TestFBXIO and (test_save_motions_with_joint_params or test_save_motions_with_model_params))'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
     "build_py",
 ] }
-install_py = { cmd = "pip install -e ." }
 doc_py = { cmd = "sphinx-build -a -E -b html pymomentum/doc build/python_api_doc", depends-on = [
-    "install_py",
+    "build_py",
 ] }
 open_doc_py = { cmd = "open build/python_api_doc/index.html", depends-on = [
     "doc_py",
@@ -340,6 +353,14 @@ install = { cmd = "cmake --build build/$PIXI_ENVIRONMENT_NAME/cpp -j --target in
     "build",
 ] }
 tracy = { cmd = "tracy-profiler.exe" }
+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" } }
+# TODO: Remove -k option, once this tests are disabled programmatically if momentum is not built with FBX support
+test_py = { cmd = "pytest pymomentum/test/*.py -k 'not (TestFBXIO and (test_save_motions_with_joint_params or test_save_motions_with_model_params))'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
+    "build_py",
+] }
 
 #==============
 # Feature: CPU
@@ -352,21 +373,21 @@ tracy = { cmd = "tracy-profiler.exe" }
 #==============
 
 [feature.py312-cuda126]
-platforms = ["linux-64"]
+platforms = ["linux-64", "win-64"]
 system-requirements = { cuda = "12" }
-build-dependencies = { cuda-toolkit = ">=12.6.3,<13", cuda-nvtx-dev = ">=12.6.3,<13" }
+build-dependencies = { cuda-toolkit = ">=12.6.3,<13", nvtx-c = ">=3.1.1" }
 dependencies = { python = "3.12.*", pytorch-gpu = ">=2.6.0,<3" }
 
 [feature.py311-cuda126]
-platforms = ["linux-64"]
+platforms = ["linux-64", "win-64"]
 system-requirements = { cuda = "12" }
-build-dependencies = { cuda-toolkit = ">=12.6.3,<13", cuda-nvtx-dev = ">=12.6.3,<13" }
+build-dependencies = { cuda-toolkit = ">=12.6.3,<13", nvtx-c = ">=3.1.1" }
 dependencies = { python = "3.11.*", pytorch-gpu = ">=2.6.0,<3" }
 
 [feature.py310-cuda126]
-platforms = ["linux-64"]
+platforms = ["linux-64", "win-64"]
 system-requirements = { cuda = "12" }
-build-dependencies = { cuda-toolkit = ">=12.6.3,<13", cuda-nvtx-dev = ">=12.6.3,<13" }
+build-dependencies = { cuda-toolkit = ">=12.6.3,<13", nvtx-c = ">=3.1.1" }
 dependencies = { python = "3.10.*", pytorch-gpu = ">=2.6.0,<3" }
 
 #==============

pixi.toml

@@ -7,22 +7,33 @@
 # 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})
+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_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
-    $ENV{CONDA_PREFIX}/lib/python${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}/site-packages/torch/lib/
+  HINTS ${libtorch_base_path}/lib/
   REQUIRED
 )
 
@@ -48,14 +59,21 @@ endif()
 # Build PyMomentum
 #===============================================================================
 
+if(WIN32)
+  set(python_deps ${Python3_LIBRARIES})
+endif()
+
 mt_library(
   NAME tensor_utility
   PYMOMENTUM_HEADERS_VARS tensor_utility_public_headers
   PYMOMENTUM_SOURCES_VARS tensor_utility_sources
   PUBLIC_INCLUDE_DIRECTORIES
+    ${Python3_INCLUDE_DIRS}
     ${TORCH_INCLUDE_DIRS}
   PUBLIC_LINK_LIBRARIES
     Eigen3::Eigen
+    pybind11::pybind11_headers
+    ${python_deps}
     ${TORCH_LIBRARIES}
   PRIVATE_LINK_LIBRARIES
     momentum
@@ -218,6 +236,7 @@ if(MOMENTUM_BUILD_TESTING)
       character_test_helpers
       tensor_ik
       tensor_utility
+      ${python_deps}
   )
 endif()