expose physx step start/finish

Author: fbxiang

include/sapien/physx/physx_default.h

@@ -21,6 +21,9 @@ class PhysxDefault {
                                  uint32_t totalAggregatePairsCapacity);
   static ::physx::PxgDynamicsMemoryConfig const &getGpuMemoryConfig();
 
+  static void setCpuWorkers(uint32_t count);
+  static uint32_t getCpuWorkers();
+
   // enable GPU simulation, may not be disabled
   static void EnableGPU();
   static bool GetGPUEnabled();

include/sapien/physx/physx_system.h

@@ -153,6 +153,10 @@ class PhysxSystemGpu : public PhysxSystem {
   getArticulationLinkComponents() const override;
 
   void step() override;
+
+  void stepStart();
+  void stepFinish();
+
   bool isGpu() const override { return true; }
 
   void gpuInit();

python/py_package/__init__.pyi

@@ -62,5 +62,5 @@ def set_cuda_tensor_backend(backend: str) -> None:
     """
 def set_log_level(level: str) -> None:
     pass
-__version__ = '3.0.0.dev20240110+ee6427de'
+__version__ = '3.0.0.dev20240110+e029ca37'
 SceneConfig = sapien.pysapien.physx.PhysxSceneConfig

python/py_package/pysapien/physx/__init__.pyi

@@ -36,8 +36,10 @@ __all__ = [
     "PhysxRigidStaticComponent",
     "PhysxSceneConfig",
     "PhysxSystem",
+    "get_cpu_workers",
     "get_default_material",
     "is_gpu_enabled",
+    "set_cpu_workers",
     "set_default_material",
     "set_gpu_memory_config"
 ]
@@ -811,6 +813,8 @@ class PhysxGpuSystem(PhysxSystem, sapien.pysapien.System):
         position `[1, 1, 1]` will be at position `[1, 1, 1] + [2, 1, 0] = [3, 2, 1]` in
         PhysX scene.
         """
+    def step_finish(self) -> None: ...
+    def step_start(self) -> None: ...
     def sync_poses_gpu_to_cpu(self) -> None: 
         """
         Warning: this function is super slow and for debug only. Download all poses from the GPU and copy to SAPIEN entities.
@@ -1273,10 +1277,14 @@ class PhysxCpuSystem(PhysxSystem, sapien.pysapien.System):
     pass
 def _enable_gpu() -> None:
     pass
+def get_cpu_workers() -> int:
+    pass
 def get_default_material() -> PhysxMaterial:
     pass
 def is_gpu_enabled() -> bool:
     pass
+def set_cpu_workers(count: int) -> None:
+    pass
 def set_default_material(static_friction: float, dynamic_friction: float, restitution: float) -> None:
     pass
 def set_gpu_memory_config(temp_buffer_capacity: int = 16777216, max_rigid_contact_count: int = 524288, max_rigid_patch_count: int = 81920, heap_capacity: int = 67108864, found_lost_pairs_capacity: int = 262144, found_lost_aggregate_pairs_capacity: int = 1024, total_aggregate_pairs_capacity: int = 1024) -> None:

python/pybind/physx.cpp

@@ -448,7 +448,10 @@ kernels will be launched to the provided stream.
 
       .def("sync_poses_gpu_to_cpu", &PhysxSystemGpu::syncPosesGpuToCpu,
            "Warning: this function is super slow and for debug only. Download all poses from the "
-           "GPU and copy to SAPIEN entities.");
+           "GPU and copy to SAPIEN entities.")
+
+      .def("step_start", &PhysxSystemGpu::stepStart)
+      .def("step_finish", &PhysxSystemGpu::stepFinish);
 
   PyPhysxMaterial
       .def(py::init<float, float, float>(), py::arg("static_friction"),
@@ -1001,7 +1004,9 @@ set some motion axes of the dynamic rigid body to be locked
            py::arg("heap_capacity") = 64 * 1024 * 1024,
            py::arg("found_lost_pairs_capacity") = 256 * 1024,
            py::arg("found_lost_aggregate_pairs_capacity") = 1024,
-           py::arg("total_aggregate_pairs_capacity") = 1024);
+           py::arg("total_aggregate_pairs_capacity") = 1024)
+      .def("set_cpu_workers", &PhysxDefault::setCpuWorkers, py::arg("count"))
+      .def("get_cpu_workers", &PhysxDefault::getCpuWorkers);
 
   ////////// end global //////////
 

src/physx/physx_default.cpp

@@ -10,6 +10,7 @@ static float gDynamicFriction{0.3};
 static float gRestitution{0.1};
 static std::weak_ptr<PhysxMaterial> gDefaultMaterial;
 static bool gGPUEnabled{false};
+static uint32_t gCpuWorkers{0};
 
 static ::physx::PxgDynamicsMemoryConfig gGpuMemoryConfig{};
 
@@ -49,6 +50,9 @@ ::physx::PxgDynamicsMemoryConfig const &PhysxDefault::getGpuMemoryConfig() {
   return gGpuMemoryConfig;
 }
 
+void PhysxDefault::setCpuWorkers(uint32_t count) { gCpuWorkers = count; }
+uint32_t PhysxDefault::getCpuWorkers() { return gCpuWorkers; }
+
 void PhysxDefault::EnableGPU() {
   if (PhysxEngine::GetIfExists()) {
     throw std::runtime_error(

src/physx/physx_system.cpp

@@ -52,7 +52,7 @@ PhysxSystemCpu::PhysxSystemCpu(PhysxSceneConfig const &config) : PhysxSystem(con
 
   sceneDesc.flags = sceneFlags;
 
-  mPxCPUDispatcher = PxDefaultCpuDispatcherCreate(0);
+  mPxCPUDispatcher = PxDefaultCpuDispatcherCreate(PhysxDefault::getCpuWorkers());
   if (!mPxCPUDispatcher) {
     throw std::runtime_error("PhysX system creation failed: failed to create CPU dispatcher");
   }
@@ -95,7 +95,7 @@ PhysxSystemGpu::PhysxSystemGpu(PhysxSceneConfig const &config) : PhysxSystem(con
 
   sceneDesc.flags = sceneFlags;
 
-  mPxCPUDispatcher = PxDefaultCpuDispatcherCreate(0);
+  mPxCPUDispatcher = PxDefaultCpuDispatcherCreate(PhysxDefault::getCpuWorkers());
   if (!mPxCPUDispatcher) {
     throw std::runtime_error("PhysX system creation failed: failed to create CPU dispatcher");
   }
@@ -190,8 +190,7 @@ std::unique_ptr<PhysxHitInfo> PhysxSystemCpu::raycast(Vec3 const &origin, Vec3 c
 
 void PhysxSystemCpu::step() {
   mPxScene->simulate(mTimestep);
-  while (!mPxScene->fetchResults(true)) {
-  }
+  mPxScene->fetchResults(true);
   for (auto c : mRigidStaticComponents) {
     c->syncPoseToEntity();
   }
@@ -210,11 +209,21 @@ void PhysxSystemGpu::step() {
 
   ++mTotalSteps;
   mPxScene->simulate(mTimestep);
-  while (!mPxScene->fetchResults(true)) {
-  }
+  mPxScene->fetchResults(true);
   // TODO: does the GPU API require fetch results?
 }
 
+void PhysxSystemGpu::stepStart() {
+  if (!mGpuInitialized) {
+    throw std::runtime_error("failed to step: gpu simulation is not initialized.");
+  }
+
+  ++mTotalSteps;
+  mPxScene->simulate(mTimestep);
+}
+
+void PhysxSystemGpu::stepFinish() { mPxScene->fetchResults(true); }
+
 std::string PhysxSystemCpu::packState() const {
   std::ostringstream ss;
   for (auto &actor : mRigidDynamicComponents) {