From 49e1ea51167febdea6bf88123b431137d8e098d9 Mon Sep 17 00:00:00 2001 From: littlemine Date: Mon, 11 Dec 2023 19:54:14 +0800 Subject: [PATCH] enhance primitiveunfuse --- projects/CUDA/utils/Primitives.cpp | 315 +++++++++++++++++------------ 1 file changed, 186 insertions(+), 129 deletions(-) diff --git a/projects/CUDA/utils/Primitives.cpp b/projects/CUDA/utils/Primitives.cpp index 7750d243e2..728ffcb5f9 100644 --- a/projects/CUDA/utils/Primitives.cpp +++ b/projects/CUDA/utils/Primitives.cpp @@ -1041,162 +1041,219 @@ ZENDEFNODE(PrimitiveFuse, { {"zs_geom"}, }); -/// @note duplicate vertices shared by multiple groups -struct PrimitiveUnfuse : INode { - void apply() override { - using namespace zs; - auto prim = get_input("prim"); - auto tag = get_input2("partition_tag"); +static std::shared_ptr unfuse_primitive(std::shared_ptr prim, std::string tag) { + using namespace zs; + constexpr auto space = execspace_e::openmp; + auto pol = omp_exec(); - constexpr auto space = execspace_e::openmp; - auto pol = omp_exec(); + auto &verts = prim->verts; + const auto &pos = verts.values; - auto &verts = prim->verts; - const auto &pos = verts.values; + auto &tris = prim->tris; + const auto &triIds = tris.values; + const bool hasTris = tris.size() > 0; - auto &tris = prim->tris; - const auto &triIds = tris.values; - const bool hasTris = tris.size() > 0; + auto &polys = prim->polys; + const auto &loops = prim->loops; + const bool hasLoops = polys.size() > 1; - auto &polys = prim->polys; - const auto &loops = prim->loops; - const bool hasLoops = polys.size() > 1; - - if ((hasTris ^ hasLoops) == 0) - throw std::runtime_error("The input mesh must either own active triangle topology or loop topology."); + if ((hasTris ^ hasLoops) == 0) + throw std::runtime_error("The input mesh must either own active triangle topology or loop topology."); - std::vector> groupsPerVertex(pos.size()); - if (hasTris) { - const auto &triGroups = tris.attr(tag); - std::vector mtxs(pos.size()); - pol(zip(triIds, triGroups), [&mtxs, &groupsPerVertex](auto tri, int groupNo) { - for (int d = 0; d != 3; ++d) { - int vi = tri[d]; - auto &mtx = mtxs[vi]; - auto &group = groupsPerVertex[vi]; - { - mtxs[vi].lock(); - group.insert(groupNo); - mtxs[vi].unlock(); - } - }; - }); - } else { - const auto &polyGroups = polys.attr(tag); - std::vector mtxs(pos.size()); - pol(zip(polys.values, polyGroups), [&mtxs, &groupsPerVertex, &loops](zeno::vec2i poly, int groupNo) { - auto st = poly[0]; - auto ed = st + poly[1]; - for (; st != ed; ++st) { - int vi = loops.values[st]; - - auto &mtx = mtxs[vi]; - auto &group = groupsPerVertex[vi]; - { - mtxs[vi].lock(); - group.insert(groupNo); - mtxs[vi].unlock(); - } + std::vector> groupsPerVertex(pos.size()); + if (hasTris) { + const auto &triGroups = tris.attr(tag); + std::vector mtxs(pos.size()); + pol(zip(triIds, triGroups), [&mtxs, &groupsPerVertex](auto tri, int groupNo) { + for (int d = 0; d != 3; ++d) { + int vi = tri[d]; + auto &mtx = mtxs[vi]; + auto &group = groupsPerVertex[vi]; + { + mtxs[vi].lock(); + group.insert(groupNo); + mtxs[vi].unlock(); } - }); + }; + }); + } else { + const auto &polyGroups = polys.attr(tag); + std::vector mtxs(pos.size()); + pol(zip(polys.values, polyGroups), [&mtxs, &groupsPerVertex, &loops](zeno::vec2i poly, int groupNo) { + auto st = poly[0]; + auto ed = st + poly[1]; + for (; st != ed; ++st) { + int vi = loops.values[st]; + + auto &mtx = mtxs[vi]; + auto &group = groupsPerVertex[vi]; + { + mtxs[vi].lock(); + group.insert(groupNo); + mtxs[vi].unlock(); + } + } + }); + } + + std::vector numGroupsPerVertex(pos.size() + 1), ptrs(pos.size() + 1); + pol(zip(numGroupsPerVertex, groupsPerVertex), [](int &num, const std::set &g) { num = g.size(); }); + exclusive_scan(pol, std::begin(numGroupsPerVertex), std::end(numGroupsPerVertex), std::begin(ptrs)); + + auto numEntries = ptrs.back(); + std::vector inds(numEntries); + pol(enumerate(groupsPerVertex), [&inds, &ptrs](int vi, const std::set &groups) { + auto st = ptrs[vi], ed = ptrs[vi + 1]; + for (auto groupNo : groups) { + inds[st++] = groupNo; // the first group does not need to change } + }); - std::vector numGroupsPerVertex(pos.size() + 1), ptrs(pos.size() + 1); - pol(zip(numGroupsPerVertex, groupsPerVertex), [](int &num, const std::set &g) { num = g.size(); }); - exclusive_scan(pol, std::begin(numGroupsPerVertex), std::end(numGroupsPerVertex), std::begin(ptrs)); + auto resPrim = std::make_shared(); - auto numEntries = ptrs.back(); - std::vector inds(numEntries); - pol(enumerate(groupsPerVertex), [&inds, &ptrs](int vi, const std::set &groups) { - auto st = ptrs[vi], ed = ptrs[vi + 1]; - for (auto groupNo : groups) { - inds[st++] = groupNo; // the first group does not need to change - } + resPrim->verts.resize(numEntries); + auto &resVerts = resPrim->verts; + auto &resPos = resVerts.values; + verts.foreach_attr([&](const auto &key, const auto &arr) { + using T = std::decay_t; + resVerts.add_attr(key); + }); + pol(range(pos.size()), [&ptrs, &verts, &resVerts, &resPos](int vi) { + auto st = ptrs[vi], ed = ptrs[vi + 1]; + for (int j = st; j != ed; ++j) + resPos[j] = verts.values[vi]; + resVerts.foreach_attr([&](const auto &key, auto &dst) { + using T = std::decay_t; + const auto &src = verts.attr(key); + for (int j = st; j != ed; ++j) + dst[j] = src[vi]; }); + }); - auto resPrim = std::make_shared(); + if (hasTris) { + const auto &triGroups = tris.attr(tag); - resPrim->verts.resize(numEntries); - auto &resVerts = resPrim->verts; - auto &resPos = resVerts.values; - verts.foreach_attr([&](const auto &key, const auto &arr) { + resPrim->tris.resize(tris.size()); + auto &resTris = resPrim->tris; + auto &resTriIds = resTris.values; + tris.foreach_attr([&](const auto &key, const auto &arr) { using T = std::decay_t; - resVerts.add_attr(key); + resTris.add_attr(key) = arr; }); - pol(range(pos.size()), [&ptrs, &verts, &resVerts, &resPos](int vi) { - auto st = ptrs[vi], ed = ptrs[vi + 1]; - for (int j = st; j != ed; ++j) - resPos[j] = verts.values[vi]; - resVerts.foreach_attr([&](const auto &key, auto &dst) { - using T = std::decay_t; - const auto &src = verts.attr(key); - for (int j = st; j != ed; ++j) - dst[j] = src[vi]; - }); + pol(range(tris.size()), [&](int f) { + auto groupNo = triGroups[f]; + for (int d = 0; d != 3; ++d) { + int vi = tris[f][d]; + int st = ptrs[vi], ed = ptrs[vi + 1]; + for (; st != ed; ++st) { + if (groupNo == inds[st]) + break; + } + resTriIds[f][d] = st; + } }); + } else { + const auto &polyGroups = polys.attr(tag); + bool uvExist = prim->uvs.size() > 0 && loops.has_attr("uvs"); - if (hasTris) { - const auto &triGroups = tris.attr(tag); + resPrim->polys.resize(polys.size()); + resPrim->loops.resize(loops.size()); + auto &resPolys = resPrim->polys; + auto &resLoops = resPrim->loops; - resPrim->tris.resize(tris.size()); - auto &resTris = resPrim->tris; - auto &resTriIds = resTris.values; - tris.foreach_attr([&](const auto &key, const auto &arr) { - using T = std::decay_t; - resTris.add_attr(key) = arr; - }); - pol(range(tris.size()), [&](int f) { - auto groupNo = triGroups[f]; - for (int d = 0; d != 3; ++d) { - int vi = tris[f][d]; - int st = ptrs[vi], ed = ptrs[vi + 1]; - for (; st != ed; ++st) { - if (groupNo == inds[st]) - break; - } - resTriIds[f][d] = st; + resPolys.values = polys.values; + polys.foreach_attr([&](const auto &key, const auto &arr) { + using T = std::decay_t; + resPolys.add_attr(key) = arr; + }); + + loops.foreach_attr([&](const auto &key, const auto &arr) { + using T = std::decay_t; + resLoops.add_attr(key) = arr; + }); + // resLoops.values = loops.values; + pol(zip(polys.values, polyGroups), [&](zeno::vec2i poly, int groupNo) { + auto st = poly[0]; + auto ed = st + poly[1]; + for (; st != ed; ++st) { + int vi = loops.values[st]; + + int l = ptrs[vi], r = ptrs[vi + 1]; + for (; l != r; ++l) { + if (groupNo == inds[l]) + break; } - }); - } else { - const auto &polyGroups = polys.attr(tag); - bool uvExist = prim->uvs.size() > 0 && loops.has_attr("uvs"); + resLoops.values[st] = l; + } + }); + } + return resPrim; +} - resPrim->polys.resize(polys.size()); - resPrim->loops.resize(loops.size()); - auto &resPolys = resPrim->polys; - auto &resLoops = resPrim->loops; +/// @note duplicate vertices shared by multiple groups +struct PrimitiveUnfuse : INode { + void apply() override { + auto prim = get_input("prim"); + auto tag = get_input2("partition_tag"); - resPolys.values = polys.values; - polys.foreach_attr([&](const auto &key, const auto &arr) { - using T = std::decay_t; - resPolys.add_attr(key) = arr; - }); + auto resPrim = unfuse_primitive(prim, tag); - loops.foreach_attr([&](const auto &key, const auto &arr) { - using T = std::decay_t; - resLoops.add_attr(key) = arr; - }); - // resLoops.values = loops.values; - pol(zip(polys.values, polyGroups), [&](zeno::vec2i poly, int groupNo) { - auto st = poly[0]; - auto ed = st + poly[1]; - for (; st != ed; ++st) { - int vi = loops.values[st]; + bool toList = get_input2("to_list"); - int l = ptrs[vi], r = ptrs[vi + 1]; - for (; l != r; ++l) { - if (groupNo == inds[l]) - break; + if (toList) { + constexpr auto space = zs::execspace_e::openmp; + auto pol = zs::omp_exec(); + + auto &verts = resPrim->verts; + const auto &pos = verts.values; + + auto &tris = resPrim->tris; + const bool hasTris = tris.size() > 0; + + auto &polys = resPrim->polys; + const bool hasLoops = polys.size() > 1; + + auto &vertGroups = resPrim->verts.add_attr(tag); + + if (hasTris) { + const auto &triIds = tris.values; + const auto &triGroups = tris.attr(tag); + + pol(zs::zip(triIds, triGroups), [&vertGroups](auto tri, int groupNo) { + for (int d = 0; d != 3; ++d) { + int vi = tri[d]; + vertGroups[vi] = groupNo; } - resLoops.values[st] = l; - } - }); + }); + } else { + const auto &loops = resPrim->loops.values; + const auto &polyGroups = polys.attr(tag); + + pol(zs::zip(polys.values, polyGroups), [&vertGroups, &loops](zeno::vec2i poly, int groupNo) { + auto st = poly[0]; + auto ed = st + poly[1]; + for (; st != ed; ++st) { + int vi = loops[st]; + vertGroups[vi] = groupNo; + } + }); + } + + auto primList = primUnmergeVerts(resPrim.get(), tag); + auto listPrim = std::make_shared(); + for (auto &primPtr : primList) { + listPrim->arr.push_back(std::move(primPtr)); + } + set_output("partitioned_prim", std::move(listPrim)); + } else { + set_output("partitioned_prim", std::move(resPrim)); } - set_output("partitioned_prim", std::move(resPrim)); } }; ZENDEFNODE(PrimitiveUnfuse, { - {{"PrimitiveObject", "prim"}, {"string", "partition_tag", "triangle_index"}}, + {{"PrimitiveObject", "prim"}, + {"string", "partition_tag", "triangle_index"}, + {"bool", "to_list", "false"}}, { {"PrimitiveObject", "partitioned_prim"}, },