Enable pymomentum on Windows

Author: jeongseok-meta

.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,19 @@ jobs:
             --build momentum/examples/hello_world/build `
             --config Release `
             --parallel
+
+  pymomentum:
+    name: py-win
+    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

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

@@ -42,6 +42,7 @@ ms-gsl = ">=4.1.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,23 @@ install = { cmd = "cmake --build build/$PIXI_ENVIRONMENT_NAME/cpp/release -j --t
     "build",
 ] }
 tracy = { cmd = "tracy-profiler" }
+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",
+] }
+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 +159,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]
@@ -226,7 +243,6 @@ 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 = """
@@ -252,7 +268,6 @@ 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 = """
@@ -340,6 +355,10 @@ 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" } }
 
 #==============
 # Feature: CPU

pixi.toml

@@ -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})
+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

pymomentum/CMakeLists.txt

@@ -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_momentum/tensor_skeleton_state.cpp

@@ -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>
 

pymomentum/tensor_utility/tensor_utility.h

@@ -27,6 +27,7 @@ cmake.args = [
     "-DMOMENTUM_USE_SYSTEM_GOOGLETEST=ON",
     "-DMOMENTUM_USE_SYSTEM_PYBIND11=OFF",
     "-DMOMENTUM_USE_SYSTEM_RERUN_CPP_SDK=ON",
+    "-DUSE_SYSTEM_NVTX=ON",
 ]
 minimum-version = "0.10"
 sdist.exclude = [
@@ -51,4 +52,5 @@ cmake.args = [
     "-DMOMENTUM_USE_SYSTEM_GOOGLETEST=ON",
     "-DMOMENTUM_USE_SYSTEM_PYBIND11=ON",
     "-DMOMENTUM_USE_SYSTEM_RERUN_CPP_SDK=ON",
+    "-DUSE_SYSTEM_NVTX=ON",
 ]