From 3316a44ad5f7ca666f4e5950265ebfcc4b08388e Mon Sep 17 00:00:00 2001
From: iaomw <iaomw@live.com>
Date: Tue, 3 Dec 2024 19:58:31 +0800
Subject: [PATCH] fix curve aux leak

---
 zenovis/src/optx/RenderEngineOptx.cpp   |  5 +++--
 zenovis/xinxinoptix/hair/optixHair.h    | 22 +++++++++++++++-------
 zenovis/xinxinoptix/optixPathTracer.cpp |  5 +----
 3 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/zenovis/src/optx/RenderEngineOptx.cpp b/zenovis/src/optx/RenderEngineOptx.cpp
index 48848fbb72..7e62e9f2cd 100644
--- a/zenovis/src/optx/RenderEngineOptx.cpp
+++ b/zenovis/src/optx/RenderEngineOptx.cpp
@@ -371,7 +371,8 @@ struct GraphicsManager {
                         }
                     }
 
-                    loadCurveGroup(points, widths, normals, strands, curveTypeEnum, mtlid);
+                    auto abcpath = ud.get2<std::string>("abcpath_0", "Default");
+                    loadCurveGroup(points, widths, normals, strands, curveTypeEnum, mtlid, abcpath);
                     return;
                 }
 
@@ -1352,7 +1353,7 @@ struct RenderEngineOptx : RenderEngine, zeno::disable_copy {
                     cachedCurvesMaterials[mtid] = { ctype };
                 }
 
-                for (auto& ele : curveGroupCache) {
+                for (auto& [key, ele] : curveGroupCache) {
 
                     auto ctype = ele->curveType;
                     auto mtlid = ele->mtlid;
diff --git a/zenovis/xinxinoptix/hair/optixHair.h b/zenovis/xinxinoptix/hair/optixHair.h
index 45da546183..e76d78de4a 100644
--- a/zenovis/xinxinoptix/hair/optixHair.h
+++ b/zenovis/xinxinoptix/hair/optixHair.h
@@ -72,6 +72,15 @@ struct HairState
     // Aux data
     CurveGroupAux aux {};
 
+    ~HairState() {
+        cudaFree( reinterpret_cast<void*>( aux.strand_u.data ) );
+        aux.strand_u.data = 0;
+        cudaFree( reinterpret_cast<void*>( aux.strand_i.data ) );
+        aux.strand_i.data = 0;
+        cudaFree( reinterpret_cast<void*>( aux.strand_info.data ) );
+        aux.strand_info.data = 0;
+    }
+
 public:
 
 void makeHairGAS(OptixDeviceContext contex);
@@ -169,11 +178,11 @@ inline void prepareHairs(OptixDeviceContext context) {
     }
 }
 
-inline std::vector<std::shared_ptr<CurveGroup>> curveGroupCache;
-inline std::vector<std::shared_ptr<HairState>> curveGroupStateCache;
+inline std::map<std::string, std::shared_ptr<CurveGroup>> curveGroupCache;
+inline std::map<std::string, std::shared_ptr<HairState>> curveGroupStateCache;
 
 inline void loadCurveGroup(std::vector<float3>& points, std::vector<float>& widths, std::vector<float3>& normals, std::vector<uint>& strands, 
-                           zeno::CurveType curveType, std::string mtlid) {
+                           zeno::CurveType curveType, std::string mtlid, std::string key) {
 
     auto cg = std::make_shared<CurveGroup>();
     cg->mtlid = mtlid; 
@@ -184,14 +193,14 @@ inline void loadCurveGroup(std::vector<float3>& points, std::vector<float>& widt
     cg->normals = std::move(normals);
     cg->strands = std::move(strands);
 
-    curveGroupCache.push_back(cg);
+    curveGroupCache[key] = cg;
 }
 
 inline void prepareCurveGroup(OptixDeviceContext context) {
 
     curveGroupStateCache.clear();
 
-    for (auto& ele : curveGroupCache) {
+    for (auto& [key, ele] : curveGroupCache) {
         auto state = std::make_shared<HairState>();
 
         state->curveGroup = ele;
@@ -200,8 +209,7 @@ inline void prepareCurveGroup(OptixDeviceContext context) {
         state->mtid = ele->mtlid;
 
         state->makeCurveGroupGAS(context);
-
-        curveGroupStateCache.push_back(state);
+        curveGroupStateCache[key] = state;
     }
 
     curveGroupCache.clear();
diff --git a/zenovis/xinxinoptix/optixPathTracer.cpp b/zenovis/xinxinoptix/optixPathTracer.cpp
index a328120fd2..68a8890f09 100644
--- a/zenovis/xinxinoptix/optixPathTracer.cpp
+++ b/zenovis/xinxinoptix/optixPathTracer.cpp
@@ -1214,12 +1214,9 @@ void updateRootIAS()
         }
     }
 
-    for (size_t i=0; i<curveGroupStateCache.size(); ++i) {
-
-        auto& ele = curveGroupStateCache[i];
+    for (auto& [key, ele] : curveGroupStateCache) {
 
         OptixInstance opinstance {};
-        
         auto shader_mark = (uint)ele->curveType + 3;
 
 		auto combinedID = ele->mtid + ":" + std::to_string(shader_mark);