From fbd6bcc59930032f7501bd6033096759c2882659 Mon Sep 17 00:00:00 2001 From: danthemighty316-jpg <273162661+danthemighty316-jpg@users.noreply.github.com> Date: Tue, 1 Sep 2026 08:36:17 +0100 Subject: [PATCH] Improve memory scaling for large 1024 generations --- include/sparse.h | 7 + src/dit.cpp | 65 +++++-- src/flow_runner.cpp | 20 +- src/remesh_dc.cpp | 15 +- src/shape_decoder.cpp | 70 +++++-- src/sparse.cpp | 417 +++++++++++++++++++++++++++++++----------- src/trellis_cli.cpp | 22 +++ src/tri_bvh.cpp | 32 +++- src/uv_bake.cpp | 147 +++++++++------ 9 files changed, 608 insertions(+), 187 deletions(-) diff --git a/include/sparse.h b/include/sparse.h index c6cebf2..2cd3497 100644 --- a/include/sparse.h +++ b/include/sparse.h @@ -23,6 +23,13 @@ ggml_tensor* sparse_submconv(ggml_context* c, const Model& m, const std::string& // SparseConvNeXtBlock3d: conv -> rowLN(affine,1e-6) -> Linear(C,4C)->SiLU->Linear(4C,C) -> + input. ggml_tensor* sparse_convnext(ggml_context* c, const Model& m, const std::string& prefix, ggml_tensor* feats, ggml_tensor* nbr, int N); +// Host-streamed ConvNeXt block for very large sparse levels. The output is identical in +// layout to sparse_convnext, but each voxel chunk is a separate short-lived GPU graph and +// only referenced neighbour feature columns are uploaded. This avoids the growing concat +// chain / full sparse-table residency that can make stage-3 shape decode request 20-30+ GB. +std::vector sparse_convnext_streamed(const Model& m, const std::string& prefix, + const std::vector& feats_in, int C, + const std::vector& nbr, int N); // SparseResBlockC2S3d up-block (channel->spatial ×2). Host-orchestrated (subdiv readback). // feats_in: [Cin*N] channel-major; returns new feats [Cout*M] + new coords (res ×2). diff --git a/src/dit.cpp b/src/dit.cpp index bc0c413..4f7ecd4 100644 --- a/src/dit.cpp +++ b/src/dit.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -65,8 +66,8 @@ static T* apply_rope(ggml_context* c, T* x, T* cos, T* sin) { // An FA padding mask [Lk_pad, Lq] (F16): 0 for real keys (< Lk_real), a large negative for the // zero-padded tail. WITHOUT it, ggml's CUDA FlashAttention folds the (zero) padded keys into the // softmax; on the >=1024-token HR flow that path NaNs a subset of queries (props, <1024 tokens, dodge -// it). WITH it the kernel masks/skips the padded KV tiles -> correct softmax, no NaN. Built once per -// flow (same N every block) and threaded into every attention; -30000 (not -inf) so 0*mask can't NaN. +// it). WITH it the kernel masks/skips the padded KV tiles -> correct softmax, no NaN. Large flows build +// this per QUERY CHUNK (not full Lq) so the mask stays bounded; -30000 (not -inf) so 0*mask can't NaN. static T* build_pad_mask(ggml_context* c, int64_t Lk_real, int64_t Lq) { const int64_t KQ = 256; const int64_t Lk_pad = ((Lk_real + KQ - 1) / KQ) * KQ; @@ -138,10 +139,52 @@ static T* sdpa(ggml_context* c, T* q, T* k, T* v, int d_model, T* mask = nullptr // ignored and the zero-padded keys are diluting the softmax (exp(0-rowmax) is only // negligible when rowmax >> 0), which shrinks every output toward zero. static const bool fa_nomask = std::getenv("TRELLIS_FA_NOMASK") != nullptr; - T* out = ggml_flash_attn_ext(c, qf, kf, vf, fa_nomask ? nullptr : mask, scale, 0.0f, 0.0f); // [hd, nh, Lq] - if (!fa_fast) ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); - out = ggml_scale(c, out, 1.0f / V_SCALE); - return ggml_reshape_2d(c, out, d_model, out->ne[2]); // [d_model, Lq] + + // IMPORTANT: the padding mask depends only on KEY index, but ggml FlashAttention + // requires it expanded to [Lk_pad, Lq_pad]. At 37,017 self-attention tokens that one + // F16 tensor is exactly 2,751,037,440 bytes -- the allocation seen in the real failure + // log. FlashAttention itself is tiled, so do the same for the mask: attention is + // independent per query, therefore query chunks are mathematically identical while + // reducing the mask from O(Lq*Lk) residency to O(nq*Lk). + const int64_t hd = qf->ne[0], Lq = qf->ne[1], nh = qf->ne[2]; + const int64_t Lk_real = k->ne[2]; + const int64_t Lk_pad = ((Lk_real + KQ_STRIDE - 1) / KQ_STRIDE) * KQ_STRIDE; + static constexpr int64_t kDefaultFaMaskChunkBytes = 256ll * 1024 * 1024; + int64_t mask_budget = kDefaultFaMaskChunkBytes; + if (const char* e = getenv("TRELLIS_FA_MASK_CHUNK_MB")) { + const int64_t mb = atoll(e); + if (mb > 0) mask_budget = mb * 1024 * 1024; + } + int64_t nq = Lq; + if (!fa_nomask) { + const int64_t bytes_per_q = std::max(1, Lk_pad * (int64_t)sizeof(uint16_t)); + nq = std::max(1, mask_budget / bytes_per_q); + // FA's mask reader works in 64-query tiles. Round ordinary chunks down to a + // multiple of 64 so only the final chunk needs padding. + if (nq >= 64 && nq < Lq) nq = (nq / 64) * 64; + if (nq > Lq) nq = Lq; + } + + T* out_all = nullptr; + for (int64_t q0 = 0; q0 < Lq; q0 += nq) { + const int64_t n = std::min(nq, Lq - q0); + T* qc = (n == Lq) + ? qf + : ggml_cont(c, ggml_view_3d(c, qf, hd, n, nh, + qf->nb[1], qf->nb[2], (size_t)q0 * qf->nb[1])); + T* cmask = nullptr; + if (!fa_nomask) { + // A caller-supplied mask is only safe to reuse when this attention was not + // query-chunked. build_dit_dense no longer creates the giant full mask. + cmask = (mask && n == Lq) ? mask : build_pad_mask(c, Lk_real, n); + } + T* o = ggml_flash_attn_ext(c, qc, kf, vf, cmask, scale, 0.0f, 0.0f); // [hd,nh,n] + if (!fa_fast) ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32); + o = ggml_scale(c, o, 1.0f / V_SCALE); + o = ggml_reshape_2d(c, o, d_model, o->ne[2]); // [d_model,n] + out_all = out_all ? ggml_concat(c, out_all, o, 1) : o; + } + return out_all; // [d_model,Lq] } // Exact SDPA, chunked over QUERIES. The whole reason FA exists here is the [Lk, Lq, nh] // score matrix -- at the HR flow that is 15104*15006*12*4 = 10.9 TB, so it cannot be @@ -275,13 +318,11 @@ ggml_tensor* build_dit_dense(ggml_context* c, const Model& m, const DiTParams& p T* mod = lin(c, m, "adaLN_modulation.1", ggml_silu(c, te));// [6*d_model] keep("t_emb_mod", mod); - // Padding masks for the two attentions, built ONCE (token counts are fixed across blocks) and - // shared by every block — they tell the CUDA FlashAttention to exclude the zero-padded key tiles - // (else the >=1024-token flow NaNs a subset of queries). self: Lk=L (the latent); cross: Lk=Lc. - T* self_mask = build_pad_mask(c, h0->ne[1], h0->ne[1]); - T* cross_mask = build_pad_mask(c, cond->ne[1], h0->ne[1]); + // FlashAttention padding masks used to be built once at full [Lk_pad,Lq_pad] size. + // That becomes a 2.75 GB single tensor at 37,017 tokens. sdpa() now creates an identical + // mask per QUERY CHUNK, so there is deliberately no full-flow mask tensor here. for (int i = 0; i < p.n_blocks; ++i) { - h = block(c, m, i, h, mod, cond, cos, sin, p, inter, self_mask, cross_mask); + h = block(c, m, i, h, mod, cond, cos, sin, p, inter, nullptr, nullptr); if (i == 0) keep("after_block0", h); if (i == 1) keep("after_block1", h); if (i == p.n_blocks - 1) keep("after_block29", h); diff --git a/src/flow_runner.cpp b/src/flow_runner.cpp index 0b5ad9e..90c961c 100644 --- a/src/flow_runner.cpp +++ b/src/flow_runner.cpp @@ -38,7 +38,12 @@ DitRunner::DitRunner(const Model& m, const DiTParams& p, int N, int n_cond, const std::vector& rcos, const std::vector& rsin) : m_(m), p_(p), N_(N), Lc_(n_cond) { const int half = p_.head_dim / 2; - size_t meta = ggml_tensor_overhead() * 16384 + ggml_graph_overhead_custom(32768, false) + (1 << 20); + // Query-chunked FlashAttention creates several small mask/FA nodes per attention instead + // of one giant quadratic mask. Give the graph generous HOST metadata headroom; this does + // not reserve an equivalent amount of VRAM. + static constexpr size_t kDitGraphNodes = 65536; + size_t meta = ggml_tensor_overhead() * kDitGraphNodes + + ggml_graph_overhead_custom(kDitGraphNodes, false) + (1 << 20); ctx_ = ggml_init({ meta, nullptr, true }); gh0_ = ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, p_.in_ch, N_); ggml_set_input(gh0_); gtf_ = ggml_new_tensor_1d(ctx_, GGML_TYPE_F32, 256); ggml_set_input(gtf_); @@ -47,12 +52,21 @@ DitRunner::DitRunner(const Model& m, const DiTParams& p, int N, int n_cond, gsin_ = ggml_new_tensor_4d(ctx_, GGML_TYPE_F32, 1, half, 1, N_); ggml_set_input(gsin_); dbg_nan_ = std::getenv("TRELLIS_DBG_NAN") != nullptr; gout_ = build_dit_dense(ctx_, m_, p_, gh0_, gtf_, gcond_, gcos_, gsin_, dbg_nan_ ? &inter_ : nullptr); - g_ = ggml_new_graph_custom(ctx_, 32768, false); + g_ = ggml_new_graph_custom(ctx_, kDitGraphNodes, false); ggml_build_forward_expand(g_, gout_); ggml_set_output(gout_); if (dbg_nan_) for (auto& [nm, t] : inter_) { ggml_build_forward_expand(g_, t); ggml_set_output(t); } alloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m_.backend)); - if (!ggml_gallocr_alloc_graph(alloc_, g_)) throw std::runtime_error("DitRunner: alloc failed"); + if (!ggml_gallocr_alloc_graph(alloc_, g_)) { + // Constructors that throw do not run DitRunner::~DitRunner(), so explicitly release + // partial allocator/context state before propagating the failure. + ggml_gallocr_free(alloc_); alloc_ = nullptr; + ggml_free(ctx_); ctx_ = nullptr; + throw std::runtime_error("DitRunner: alloc failed"); + } + if (getenv("TRELLIS_DBG_ALLOC")) + fprintf(stderr, " [dit-alloc] N=%d nodes=%d gallocr buffer = %.2f GB\n", + N_, ggml_graph_n_nodes(g_), ggml_gallocr_get_buffer_size(alloc_, 0) / 1e9); rcos_ = rcos; rsin_ = rsin; // keep; re-upload each forward (gallocr reuses input buffers across runs) } diff --git a/src/remesh_dc.cpp b/src/remesh_dc.cpp index f43dec0..b3440ec 100644 --- a/src/remesh_dc.cpp +++ b/src/remesh_dc.cpp @@ -76,7 +76,20 @@ Mesh remesh_narrow_band_dc(const float* iverts, int64_t iV, const int32_t* iface { const int F = (int)iF; std::vector> parts; - const int nt = std::max(1u, std::thread::hardware_concurrency()); + const int hw = (int)std::max(1u, std::thread::hardware_concurrency()); + // Each worker previously allocated a full res^3 candidate bitset. At + // res=1024 that is 128 MiB per worker, so a 32-thread CPU consumed ~4 GiB + // here before any geometry/BVH memory. Cap only the replication memory; + // the OR result and therefore the remesh are bit-for-bit equivalent. + const size_t bytes_per_part = cand.size() * sizeof(uint64_t); + const size_t parts_budget = (size_t)1024 * 1024 * 1024; // 1 GiB + const int mem_workers = bytes_per_part ? (int)std::max(1, parts_budget / bytes_per_part) : hw; + const int nt = std::max(1, std::min(hw, mem_workers)); + if (nt < hw) { + printf(" [remesh-mem] candidate bitset %.1f MiB/worker, workers %d->%d\n", + bytes_per_part / (1024.0*1024.0), hw, nt); + fflush(stdout); + } parts.assign(nt, {}); std::vector ts; const int chunk = (F + nt - 1) / nt; diff --git a/src/shape_decoder.cpp b/src/shape_decoder.cpp index 2aea89f..c1d4a20 100644 --- a/src/shape_decoder.cpp +++ b/src/shape_decoder.cpp @@ -6,6 +6,7 @@ #include "ggml-alloc.h" #include +#include #include #include @@ -37,12 +38,20 @@ static std::vector run1(const Model& m, ggml_context* c, T* out, ggml_cgraph* g = ggml_new_graph_custom(c, kGraphNodes, false); ggml_build_forward_expand(g, out); ggml_gallocr_t a = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m.backend)); - if (!ggml_gallocr_alloc_graph(a, g)) throw std::runtime_error("shape_dec alloc"); + if (!ggml_gallocr_alloc_graph(a, g)) { + // A failed reserve can still leave allocator/backend bookkeeping behind. Free it before + // the caller retries the same stage through the streamed ConvNeXt fallback. + ggml_gallocr_free(a); + throw std::runtime_error("shape_dec alloc"); + } if (getenv("TRELLIS_DBG_ALLOC")) fprintf(stderr, " [stage-alloc] nodes=%d gallocr buffer = %.2f GB\n", ggml_graph_n_nodes(g), ggml_gallocr_get_buffer_size(a, 0) / 1e9); for (auto& [t, d] : ins) ggml_backend_tensor_set(t, d, 0, ggml_nbytes(t)); - if (ggml_backend_graph_compute(m.backend, g) != GGML_STATUS_SUCCESS) throw std::runtime_error("shape_dec compute"); + if (ggml_backend_graph_compute(m.backend, g) != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(a); + throw std::runtime_error("shape_dec compute"); + } mem_probe("run1 computed (pre-readback)"); std::vector r = tensor_to_f32(out); ggml_gallocr_free(a); @@ -94,15 +103,54 @@ static std::vector decode_unet(const Model& m, const std::vector& if (getenv("TRELLIS_DBG_MEM")) fprintf(stderr, " == stage %d: C=%d nblk=%d N=%d | host h=%.2f GB nbr=%.2f GB\n", si, st.C, st.nblk, N, (double)h.size() * 4 / 1e9, (double)nbr.size() * 4 / 1e9); - { // ConvNeXt stage - ggml_context* c = mkctx(); - T* gh = ggml_new_tensor_2d(c, GGML_TYPE_F32, st.C, N); ggml_set_input(gh); - T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, N, 27); ggml_set_input(gn); - T* x = gh; - for (int j = 0; j < st.nblk; ++j) - x = sparse_convnext(c, m, std::string("blocks.") + st.s + "." + std::to_string(j), x, gn, N); - h = run1(m, c, x, { {gh, h.data()}, {gn, nbr.data()} }); - ggml_free(c); + { // ConvNeXt stage. Once the full stage input itself reaches ~2 GiB, the old + // sparse_convnext concat chain multiplies that residency several-fold (the real + // stage-3 failure reserved 31.36 GB). Go straight to true chunk-local streaming + // for that case instead of deliberately provoking an OOM first. Smaller stages + // keep the fast whole-stage graph, with an allocation-failure fallback as a safety net. + size_t stream_at = 2ull * 1024 * 1024 * 1024; + if (const char* e = getenv("TRELLIS_CONVNEXT_STREAM_GB")) { + const double gb = atof(e); + if (gb > 0.0) stream_at = (size_t)(gb * 1024.0 * 1024.0 * 1024.0); + } + const size_t feature_bytes = (size_t)st.C * (size_t)N * sizeof(float); + const bool force_stream = getenv("TRELLIS_CONVNEXT_STREAM_ALL") != nullptr; + bool stream_stage = force_stream || feature_bytes >= stream_at; + + if (!stream_stage) { + // The host vector `h` is only assigned after a successful run1(), so it remains + // the stage input if the fast-path allocation fails. + ggml_context* c = mkctx(); + try { + T* gh = ggml_new_tensor_2d(c, GGML_TYPE_F32, st.C, N); ggml_set_input(gh); + T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, N, 27); ggml_set_input(gn); + T* x = gh; + for (int j = 0; j < st.nblk; ++j) + x = sparse_convnext(c, m, std::string("blocks.") + st.s + "." + std::to_string(j), x, gn, N); + std::vector fast = run1(m, c, x, { {gh, h.data()}, {gn, nbr.data()} }); + ggml_free(c); c = nullptr; + h = std::move(fast); + } catch (const std::runtime_error& e) { + if (c) ggml_free(c); + if (std::string(e.what()) != "shape_dec alloc") throw; + stream_stage = true; + fprintf(stderr, + " [shape-dec-stream] stage %d C=%d N=%d whole-stage alloc failed; " + "retrying streamed\n", si, st.C, N); + } + } else if (getenv("TRELLIS_DBG_ALLOC")) { + fprintf(stderr, + " [shape-dec-stream] stage %d C=%d N=%d feature-table=%.2f GiB " + ">= %.2f GiB; streaming\n", + si, st.C, N, feature_bytes / 1073741824.0, stream_at / 1073741824.0); + } + + if (stream_stage) { + for (int j = 0; j < st.nblk; ++j) { + const std::string bp = std::string("blocks.") + st.s + "." + std::to_string(j); + h = sparse_convnext_streamed(m, bp, h, st.C, nbr, N); + } + } } mem_probe("after ConvNeXt stage"); const std::vector* ext = guide_subs ? &(*guide_subs)[si] : nullptr; diff --git a/src/sparse.cpp b/src/sparse.cpp index 1a0b57b..e4cf20f 100644 --- a/src/sparse.cpp +++ b/src/sparse.cpp @@ -139,9 +139,9 @@ ggml_tensor* sparse_convnext(ggml_context* c, const Model& m, const std::string& // Run a graph with given inputs (name->host data), return one named output to host. namespace { -// c2s builds conv1 + conv2 into ONE graph, and each conv is 27 taps x mul_mat_rows' 1M-row -// chunks -- at res-1024 the post-subdiv conv2 alone is ~16 chunks. Node/tensor meta is host -// RAM for structs only (~370 B each), so budget generously rather than risk GGML_ASSERT. +// C2S streaming uses many short-lived GraphRun instances. Keep the metadata ceiling generous: +// node/tensor metadata is host RAM only (~370 B each), while GPU workspace is owned and freed +// independently by each GraphRun's gallocr. static constexpr size_t kGraphNodes = 65536; struct GraphRun { const Model& m; ggml_context* c; ggml_gallocr_t alloc = nullptr; @@ -170,6 +170,125 @@ struct GraphRun { }; } // anon +// Host-streamed SparseConvNeXt block. This is deliberately separate from the graph-builder +// sparse_convnext() above: the latter is excellent for ordinary sparse levels, but for very +// large stage-3 levels its chunk outputs are joined by a growing ggml_concat chain. The +// allocator then has to reserve every prefix of that chain and a single block can become a +// multi-GB monolith. Here every output chunk is an independent GraphRun, stitched on the host. +// As with compact C2S, only feature columns actually referenced by this chunk's 27-neighbour +// table are uploaded, so the full [C,N] sparse feature table never has to reside on the GPU. +std::vector sparse_convnext_streamed(const Model& m, const std::string& prefix, + const std::vector& feats_in, int C, + const std::vector& nbr, int N) { + if (N <= 0 || C <= 0 || feats_in.size() != (size_t)C * N) + throw std::runtime_error("shape_dec: invalid streamed ConvNeXt input in " + prefix); + if (nbr.size() != (size_t)27 * N) + throw std::runtime_error("shape_dec: invalid streamed ConvNeXt neighbour table in " + prefix); + + struct CompactNeighbors { + std::vector table; // tap-major local indices, sentinel = globals.size() + std::vector globals; // local column -> source column in feats_in + }; + auto compact_neighbors = [](const std::vector& full, int full_n, + int64_t r0, int nr, int sentinel) { + CompactNeighbors out; + out.table.resize((size_t)27 * nr); + std::unordered_map remap; + remap.reserve((size_t)nr * 2 + 1); + for (int t = 0; t < 27; ++t) { + for (int j = 0; j < nr; ++j) { + const int32_t g = full[(size_t)t * full_n + (size_t)r0 + j]; + if (g == sentinel) { + out.table[(size_t)t * nr + j] = -1; + continue; + } + auto it = remap.find(g); + if (it == remap.end()) { + const int32_t li = (int32_t)out.globals.size(); + out.globals.push_back(g); + remap.emplace(g, li); + out.table[(size_t)t * nr + j] = li; + } else { + out.table[(size_t)t * nr + j] = it->second; + } + } + } + const int32_t local_sentinel = (int32_t)out.globals.size(); + for (int32_t& x : out.table) if (x < 0) x = local_sentinel; + return out; + }; + auto gather_columns = [](const std::vector& full, int ch, + const std::vector& globals) { + std::vector local((size_t)ch * (globals.size() + 1), 0.0f); + for (size_t li = 0; li < globals.size(); ++li) { + const size_t src = (size_t)ch * (size_t)globals[li]; + std::copy_n(full.data() + src, ch, local.data() + (size_t)ch * li); + } + return local; + }; + + static constexpr int64_t kDefaultConvNextChunkBytes = 256ll * 1024 * 1024; + static constexpr int64_t kMaxConvNextStreamVoxels = 65536; + int64_t budget = kDefaultConvNextChunkBytes; + if (const char* e = getenv("TRELLIS_CONVNEXT_CHUNK_MB")) { + const int64_t mb = atoll(e); + if (mb > 0) budget = mb * 1024 * 1024; + } + // The MLP widens C -> 4C and is the widest voxel-local activation. Keep that wide + // activation near the requested budget, then cap output rows just like compact C2S. + const int64_t per_vox = std::max(1, 4ll * C * (int64_t)sizeof(float)); + int64_t chunk = std::max(1, budget / per_vox); + if (chunk > kMaxConvNextStreamVoxels) chunk = kMaxConvNextStreamVoxels; + if (chunk > N) chunk = N; + const int64_t nchunks = (N + chunk - 1) / chunk; + + if (getenv("TRELLIS_DBG_ALLOC")) + fprintf(stderr, " [convnext-stream] %s C=%d N=%d chunks=%lldx%lld budget=%lld MB\n", + prefix.c_str(), C, N, (long long)nchunks, (long long)chunk, + (long long)(budget / (1024 * 1024))); + + std::vector out((size_t)C * N); + for (int64_t r0 = 0; r0 < N; r0 += chunk) { + const int nr = (int)std::min(chunk, N - r0); + CompactNeighbors cn = compact_neighbors(nbr, N, r0, nr, N); + std::vector local = gather_columns(feats_in, C, cn.globals); + + GraphRun gr(m); + ggml_context* c = gr.c; + const int nsrc = (int)cn.globals.size(); + T* gh = ggml_new_tensor_2d(c, GGML_TYPE_F32, C, nsrc + 1); ggml_set_input(gh); + T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, nr, 27); ggml_set_input(gn); + T* gres = ggml_new_tensor_2d(c, GGML_TYPE_F32, C, nr); ggml_set_input(gres); + + T* Wc = w32(c, m.get(prefix + ".conv.weight")); + T* h = submconv_range(c, m, prefix + ".conv", gh, Wc, gn, nr, 0, nr); + h = ggml_norm(c, h, 1e-6f); + h = ggml_add(c, ggml_mul(c, h, m.get(prefix + ".norm.weight")), + m.get(prefix + ".norm.bias")); + h = ggml_add(c, mul_mat_rows(c, m.get(prefix + ".mlp.0.weight"), h), + m.get(prefix + ".mlp.0.bias")); + h = ggml_silu(c, h); + h = ggml_add(c, mul_mat_rows(c, m.get(prefix + ".mlp.2.weight"), h), + m.get(prefix + ".mlp.2.bias")); + h = ggml_add(c, h, gres); + + std::vector hv; + try { + hv = gr.run(h, { {gh, local.data()}, {gn, cn.table.data()}, + {gres, feats_in.data() + (size_t)C * r0} }); + } catch (const std::runtime_error& e) { + // GraphRun's generic label is c2s because it was originally introduced for C2S. + // Reclassify here so the caller/queue reports the subsystem that actually failed. + const std::string what = e.what(); + if (what == "c2s: alloc failed") throw std::runtime_error("shape_dec alloc"); + if (what == "c2s: compute failed") throw std::runtime_error("shape_dec compute"); + throw; + } + std::copy(hv.begin(), hv.end(), out.begin() + (size_t)C * r0); + } + return out; +} + C2SResult sparse_c2s(const Model& m, const std::string& prefix, const std::vector& feats_in, int Cin, const std::vector>& coords, int Cout, @@ -189,12 +308,30 @@ C2SResult sparse_c2s(const Model& m, const std::string& prefix, std::vector subdiv; if (!ext_subdiv) { auto predict_subdiv = [&]() { - GraphRun gr1(m); - ggml_context* c = gr1.c; - T* gf = ggml_new_tensor_2d(c, GGML_TYPE_F32, Cin, N); ggml_set_input(gf); - T* sd = ggml_add(c, mul_mat_rows(c, m.get(prefix + ".to_subdiv.weight"), gf), - m.get(prefix + ".to_subdiv.bias")); - return gr1.run(sd, { {gf, feats_in.data()} }); // [8, N] -- small + // to_subdiv is a voxel-local linear projection [Cin,N] -> [8,N]. Keeping the + // whole F32 input tensor on-device defeats C2S streaming for very dense shape + // decoder stages (e.g. Cin=128,N~4.5M is >2 GiB before any workspace). Stream + // independent voxel ranges and stitch the tiny 8-channel logits on the host. + static constexpr int64_t kSubdivStreamVoxels = 65536; + std::vector out((size_t)8 * N); + if (N > kSubdivStreamVoxels) + fprintf(stderr, + " [c2s-subdiv-stream] %s N=%d chunks=%lldx%lld\n", + prefix.c_str(), N, + (long long)((N + kSubdivStreamVoxels - 1) / kSubdivStreamVoxels), + (long long)kSubdivStreamVoxels); + + for (int64_t r0 = 0; r0 < N; r0 += kSubdivStreamVoxels) { + const int nr = (int)std::min(kSubdivStreamVoxels, (int64_t)N - r0); + GraphRun gr1(m); + ggml_context* c = gr1.c; + T* gf = ggml_new_tensor_2d(c, GGML_TYPE_F32, Cin, nr); ggml_set_input(gf); + T* sd = ggml_add(c, ggml_mul_mat(c, m.get(prefix + ".to_subdiv.weight"), gf), + m.get(prefix + ".to_subdiv.bias")); + std::vector sv = gr1.run(sd, { {gf, feats_in.data() + (size_t)Cin * r0} }); + std::copy(sv.begin(), sv.end(), out.begin() + (size_t)8 * r0); + } + return out; }; constexpr int kAttempts = 3; for (int attempt = 0; attempt < kAttempts; ++attempt) { @@ -252,111 +389,185 @@ C2SResult sparse_c2s(const Model& m, const std::string& prefix, prefix.c_str(), N, M, (double)nnbr.size() * 4 / 1e9, (double)feats_in.size() * 4 / 1e9, (double)Cout * M * 4 / 1e9); - // conv1 widens to Cout*8, so the whole [Cout*8, N] is the stage's largest tensor (2.5 GB at - // res-1024 stage 3). Materialising it whole and gathering afterwards keeps it live next to - // conv2's working set -- measured 10.1 GB for the stage. Chunking conv1 over input voxels - // and gathering each chunk on the spot bounds it to [Cout*8, nr] instead. - const int64_t per_vox = (int64_t)Cout * 8 * 4; - int64_t budget = kBlockChunkBytes; - if (const char* e = getenv("TRELLIS_C2S_CHUNK_MB")) budget = atoll(e) * 1024 * 1024; - int64_t chunk = std::max(1, budget / std::max(per_vox, 1)); - constexpr int64_t kMaxChunks = 24; // node-budget floor, as in sparse_convnext - if (chunk * kMaxChunks < N) chunk = (N + kMaxChunks - 1) / kMaxChunks; - if (chunk >= N) chunk = N; + // ---- streamed C2S --------------------------------------------------------- + // The older implementation chunked conv1/conv2 logically but still built every chunk + // into ONE ggml graph. gallocr therefore reserved the aggregate lifetime of the whole + // graph, which defeats the memory cap and can request 10-20+ GB at res-1024. + // + // This path makes the chunk boundary a real allocation boundary: each normalization, + // conv1 and conv2 chunk gets its own GraphRun. The GraphRun is destroyed before the next + // chunk starts, so its gallocr buffer is released. Intermediate C2S features are stitched + // on the host. This trades some PCIe traffic for a bounded GPU working set -- exactly what + // the 16 GB Windows/Vulkan path needs. + static constexpr int64_t kDefaultC2SChunkBytes = 256ll * 1024 * 1024; + int64_t budget = kDefaultC2SChunkBytes; + if (const char* e = getenv("TRELLIS_C2S_CHUNK_MB")) { + const int64_t mb = atoll(e); + if (mb > 0) budget = mb * 1024 * 1024; + } - // Per-chunk gather indices are chunk-local (the [Cout, 8*nr] reshape restarts at 0), so - // rebase once on the host: gloc[m] = gidx[m] - 8*r0 for m in this chunk's range. - std::vector gloc(M); - for (int64_t r0 = 0; r0 < N; r0 += chunk) { - const int64_t r1 = std::min(r0 + chunk, N); - for (int32_t mm = mstart[r0]; mm < mstart[r1]; ++mm) gloc[mm] = gidx[mm] - (int32_t)(8 * r0); + auto chunk_count = [](int64_t total, int64_t chunk_sz) -> int64_t { + return (total + chunk_sz - 1) / chunk_sz; + }; + struct CompactNeighbors { + std::vector table; // tap-major local indices, sentinel = globals.size() + std::vector globals; // local column -> source column in the full host table + }; + auto compact_neighbors = [](const std::vector& full, int full_n, + int64_t r0, int nr, int sentinel) { + CompactNeighbors out; + out.table.resize((size_t)27 * nr); + // A streamed graph must not upload the full [C,N] feature table just because its + // neighbours are global indices. Compact the referenced columns for this output chunk + // and remap the neighbour table to that compact table. With a 65k output chunk this + // stays small even when the full post-subdivision stage has 10M+ voxels. + std::unordered_map remap; + remap.reserve((size_t)nr * 2 + 1); + for (int t = 0; t < 27; ++t) { + for (int j = 0; j < nr; ++j) { + const int32_t g = full[(size_t)t * full_n + (size_t)r0 + j]; + if (g == sentinel) { + out.table[(size_t)t * nr + j] = -1; + continue; + } + auto it = remap.find(g); + if (it == remap.end()) { + const int32_t li = (int32_t)out.globals.size(); + out.globals.push_back(g); + remap.emplace(g, li); + out.table[(size_t)t * nr + j] = li; + } else { + out.table[(size_t)t * nr + j] = it->second; + } + } + } + const int32_t local_sentinel = (int32_t)out.globals.size(); + for (int32_t& x : out.table) if (x < 0) x = local_sentinel; + return out; + }; + auto gather_columns = [](const std::vector& full, int C, + const std::vector& globals) { + std::vector local((size_t)C * (globals.size() + 1), 0.0f); + for (size_t li = 0; li < globals.size(); ++li) { + const size_t src = (size_t)C * (size_t)globals[li]; + std::copy_n(full.data() + src, C, local.data() + (size_t)C * li); + } + return local; + }; + + // norm1 + affine + SiLU are voxel-local. Do them in independent row chunks and keep the + // result on the host with one extra zero sentinel row. This also avoids submconv_pad's + // full-size device concat in every conv1 chunk. + const int64_t norm_bytes_per_vox = std::max(1, 4ll * Cin * (int64_t)sizeof(float)); + int64_t norm_chunk = std::max(1, budget / norm_bytes_per_vox); + if (norm_chunk > N) norm_chunk = N; + + std::vector hnorm((size_t)Cin * (N + 1), 0.0f); // column N is the sentinel + for (int64_t r0 = 0; r0 < N; r0 += norm_chunk) { + const int nr = (int)std::min(norm_chunk, N - r0); + GraphRun gr(m); + ggml_context* c = gr.c; + T* gf = ggml_new_tensor_2d(c, GGML_TYPE_F32, Cin, nr); ggml_set_input(gf); + T* h = ggml_norm(c, gf, 1e-6f); + h = ggml_add(c, ggml_mul(c, h, m.get(prefix + ".norm1.weight")), + m.get(prefix + ".norm1.bias")); + h = ggml_silu(c, h); + std::vector hv = gr.run(h, { {gf, feats_in.data() + (size_t)Cin * r0} }); + std::copy(hv.begin(), hv.end(), hnorm.begin() + (size_t)Cin * r0); } - // ---- graph 2: conv1 -> SparseChannel2Spatial -> conv2 + skip, entirely on device ---- - GraphRun gr2(m); - ggml_context* c = gr2.c; - T* gf = ggml_new_tensor_2d(c, GGML_TYPE_F32, Cin, N); ggml_set_input(gf); - T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, N, 27); ggml_set_input(gn); - T* gi = ggml_new_tensor_1d(c, GGML_TYPE_I32, M); ggml_set_input(gi); - T* gl = ggml_new_tensor_1d(c, GGML_TYPE_I32, M); ggml_set_input(gl); - T* gn2 = ggml_new_tensor_2d(c, GGML_TYPE_I32, M, 27); ggml_set_input(gn2); - - T* h = ggml_norm(c, gf, 1e-6f); - h = ggml_add(c, ggml_mul(c, h, m.get(prefix + ".norm1.weight")), m.get(prefix + ".norm1.bias")); - h = ggml_silu(c, h); - - // SparseChannel2Spatial(2). [Cout*8, nr] is contiguous, so octant o's slice of voxel i -- - // channels (o*Cout .. o*Cout+Cout), at offset k + Cout*(o + 8*i) -- is exactly column - // (o + 8*i) of the free [Cout, 8*nr] reshape. The subdivision is then just a gather of the - // surviving columns: no host copy, no re-upload. - // Chunks are written into ONE [Cout, M+1] buffer rather than concat-chained: every - // ggml_concat allocates a new full-size tensor beside the old one, so a chain peaks at - // ~2x its result -- 16.7 GB for the cottage's [64, 32.7M]. ggml_pad seeds the buffer from - // chunk 0 and zero-fills the rest, so column M -- conv2's sentinel row, which absent - // neighbours index -- is already zero; norm2/silu map an all-zero column to zero - // (norm is (0-0)/sqrt(0+eps)=0, silu(0)=0), so it survives and submconv_pad's extra full - // copy of hn is gone too. - // ggml_cpy, not ggml_set_2d_inplace: SET keeps its byte offset in int32 op_params and - // asserts offset < 1<<30, which these multi-GB tensors blow past. A view carries a 64-bit - // pointer instead. Nothing reads the writes, so they are rooted explicitly via run()'s - // `roots`, expanded ahead of the consumers that follow them in the node array. - T* W1 = w32(c, m.get(prefix + ".conv1.weight")); - T* fz = submconv_pad(c, h, W1->ne[0]); // built once, shared by every chunk - std::vector roots; - T* hraw = nullptr; + // conv1 widens to Cout*8. Each input-voxel range maps to a contiguous output range + // [mstart[r0], mstart[r1]), so we can gather the surviving octants immediately, apply + // norm2+SiLU immediately, read that chunk back, and discard the graph before moving on. + const int64_t per_vox = std::max(1, (int64_t)Cout * 8 * (int64_t)sizeof(float)); + int64_t chunk = std::max(1, budget / per_vox); + // submconv has 27 gathers/matmuls; the output tensor alone is not a useful estimate of + // its real Vulkan working set. Keep each streamed conv graph comfortably below the + // multi-GB allocations seen on 16 GB cards. + static constexpr int64_t kMaxC2SStreamVoxels = 65536; + if (chunk > kMaxC2SStreamVoxels) chunk = kMaxC2SStreamVoxels; + if (chunk > N) chunk = N; + + std::vector hn((size_t)Cout * (M + 1), 0.0f); // host table; column M is sentinel for (int64_t r0 = 0; r0 < N; r0 += chunk) { const int64_t r1 = std::min(r0 + chunk, N); - const int32_t m0 = mstart[r0], m1 = mstart[r1]; - if (m1 == m0) continue; // no octant of this chunk survived - T* cvc = submconv_range(c, m, prefix + ".conv1", fz, W1, gn, N, (int)r0, (int)(r1 - r0)); - T* idx = ggml_cont(c, ggml_view_1d(c, gl, m1 - m0, (size_t)m0 * ggml_element_size(gl))); - T* hc = ggml_get_rows(c, ggml_reshape_2d(c, cvc, Cout, 8 * (r1 - r0)), idx); // [Cout, m1-m0] - if (!hraw) { hraw = ggml_pad(c, hc, 0, (int)(M + 1 - (m1 - m0)), 0, 0); continue; } // [Cout, M+1] - roots.push_back(ggml_cpy(c, hc, ggml_view_2d(c, hraw, Cout, m1 - m0, - hraw->nb[1], (size_t)m0 * hraw->nb[1]))); + const int nr = (int)(r1 - r0); + const int32_t m0 = mstart[(size_t)r0], m1 = mstart[(size_t)r1]; + if (m1 == m0) continue; + const int mc = m1 - m0; + + CompactNeighbors cn = compact_neighbors(nbr, N, r0, nr, N); + std::vector hnorm_local = gather_columns(hnorm, Cin, cn.globals); + std::vector idx_local((size_t)mc); + for (int j = 0; j < mc; ++j) + idx_local[(size_t)j] = gidx[(size_t)m0 + j] - (int32_t)(8 * r0); + + GraphRun gr(m); + ggml_context* c = gr.c; + const int nsrc = (int)cn.globals.size(); + T* gh = ggml_new_tensor_2d(c, GGML_TYPE_F32, Cin, nsrc + 1); ggml_set_input(gh); + T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, nr, 27); ggml_set_input(gn); + T* gi = ggml_new_tensor_1d(c, GGML_TYPE_I32, mc); ggml_set_input(gi); + + T* W1 = w32(c, m.get(prefix + ".conv1.weight")); + T* cvc = submconv_range(c, m, prefix + ".conv1", gh, W1, gn, nr, 0, nr); + T* hc = ggml_get_rows(c, ggml_reshape_2d(c, cvc, Cout, (int64_t)8 * nr), gi); + hc = ggml_silu(c, ggml_norm(c, hc, 1e-6f)); + + std::vector hv = gr.run(hc, { {gh, hnorm_local.data()}, {gn, cn.table.data()}, + {gi, idx_local.data()} }); + std::copy(hv.begin(), hv.end(), hn.begin() + (size_t)Cout * m0); } - T* hn = ggml_silu(c, ggml_norm(c, hraw, 1e-6f)); // norm2, no affine - - // conv2 runs at the POST-subdivision count M, which is where a dense object lives: the - // cottage hits M=32.7M, and unchunked each of the 27 taps builds full [Cout, M] gather / - // matmul / accumulator tensors, while mul_mat_rows' own 1M-row split concats 33 pieces - // into a chain that peaks near 2x the output. Chunking over OUTPUT voxels fixes both -- - // the gather still reads the whole (padded) hn, since a neighbour can be any voxel, but - // the output is per-voxel so a range needs no halo. Keeping chunks <= mul_mat_rows' - // 1M-row threshold also stops it splitting internally, removing that concat chain. - // skip rides along per chunk: x channel2spatial'd by the same gather ([Cin,N] -> [K,M]), - // then repeat_interleave(R). [1,K,nr] -> repeat -> [R,K,nr] lays element (r,k,m) at - // r + R*k + R*K*m, so the [Cout,nr] reshape maps channel k*R+r -> k: interleave, not tile. - T* W2 = w32(c, m.get(prefix + ".conv2.weight")); - T* fz2 = hn; // already [Cout, M+1]: pre-padded above - T* xs = ggml_get_rows(c, ggml_reshape_2d(c, gf, K, (int64_t)8 * N), gi); // [K, M] - - const int64_t per_vox2 = 3 * (int64_t)Cout * 4; // acc + gather + matmul, live per tap - int64_t chunk2 = std::max(1, budget / std::max(per_vox2, 1)); + + // conv2 is streamed over POST-subdivision voxels. Neighbours are global on the host, but + // each chunk compacts just the referenced feature columns and remaps its neighbour table, + // so the GPU never receives the full [Cout,M] table. There is deliberately no full + // [Cout,M] device input or output: each chunk is read back and stitched into `outv`. + const int64_t per_vox2 = std::max(1, 3ll * Cout * (int64_t)sizeof(float)); + int64_t chunk2 = std::max(1, budget / per_vox2); + if (chunk2 > kMaxC2SStreamVoxels) chunk2 = kMaxC2SStreamVoxels; if (chunk2 > kMulMatRowChunk) chunk2 = kMulMatRowChunk; - constexpr int64_t kMaxChunks2 = 64; // ~135 nodes/chunk, well under kGraphNodes - if (chunk2 * kMaxChunks2 < M) chunk2 = (M + kMaxChunks2 - 1) / kMaxChunks2; - if (chunk2 >= M) chunk2 = M; + if (chunk2 > M) chunk2 = M; - T* out = nullptr; + if (getenv("TRELLIS_DBG_ALLOC")) + fprintf(stderr, + " [c2s-stream] %s budget=%lld MB norm=%lldx%lld conv1=%lldx%lld conv2=%lldx%lld\n", + prefix.c_str(), (long long)(budget / (1024 * 1024)), + (long long)chunk_count(N, norm_chunk), (long long)norm_chunk, + (long long)chunk_count(N, chunk), (long long)chunk, + (long long)chunk_count(M, chunk2), (long long)chunk2); + + std::vector outv((size_t)Cout * M); for (int64_t m0 = 0; m0 < M; m0 += chunk2) { - const int64_t nr2 = std::min(chunk2, M - m0); - T* o = submconv_range(c, m, prefix + ".conv2", fz2, W2, gn2, M, (int)m0, (int)nr2); - T* xr = (m0 == 0 && nr2 == M) ? xs - : ggml_cont(c, ggml_view_2d(c, xs, K, nr2, xs->nb[1], (size_t)m0 * xs->nb[1])); - T* skr = ggml_repeat_4d(c, ggml_reshape_3d(c, xr, 1, K, nr2), R, K, nr2, 1); - o = ggml_add(c, o, ggml_reshape_2d(c, skr, Cout, nr2)); - // written into one [Cout, M] buffer, as for hraw above -- a concat chain here would - // again peak at 2x the 8.4 GB output. Every column is covered by some chunk, so the - // pad's zero fill is overwritten and only sizes the buffer. - if (!out) { out = ggml_pad(c, o, 0, (int)(M - nr2), 0, 0); continue; } // [Cout, M] - roots.push_back(ggml_cpy(c, o, ggml_view_2d(c, out, Cout, nr2, - out->nb[1], (size_t)m0 * out->nb[1]))); + const int nr2 = (int)std::min(chunk2, M - m0); + CompactNeighbors cn = compact_neighbors(nnbr, M, m0, nr2, M); + std::vector hn_local = gather_columns(hn, Cout, cn.globals); + + GraphRun gr(m); + ggml_context* c = gr.c; + const int nsrc = (int)cn.globals.size(); + T* gh = ggml_new_tensor_2d(c, GGML_TYPE_F32, Cout, nsrc + 1); ggml_set_input(gh); + T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, nr2, 27); ggml_set_input(gn); + T* W2 = w32(c, m.get(prefix + ".conv2.weight")); + T* o = submconv_range(c, m, prefix + ".conv2", gh, W2, gn, nr2, 0, nr2); + + std::vector ov = gr.run(o, { {gh, hn_local.data()}, {gn, cn.table.data()} }); + + // Skip path on host: channel2spatial(raw x) followed by repeat_interleave(R). + // gidx[m] is the selected column in reshape(raw_x, [K, 8N]). + for (int j = 0; j < nr2; ++j) { + const int64_t mm = m0 + j; + const int32_t src_col = gidx[(size_t)mm]; + for (int k = 0; k < K; ++k) { + const float sv = feats_in[(size_t)K * src_col + k]; + const int c0 = k * R; + for (int r = 0; r < R; ++r) + ov[(size_t)Cout * j + c0 + r] += sv; + } + } + std::copy(ov.begin(), ov.end(), outv.begin() + (size_t)Cout * m0); } - std::vector outv = gr2.run(out, { {gf, feats_in.data()}, {gn, nbr.data()}, - {gi, gidx.data()}, {gl, gloc.data()}, {gn2, nnbr.data()} }, - roots); return { std::move(outv), std::move(nc), Cout, std::move(mask_used) }; } diff --git a/src/trellis_cli.cpp b/src/trellis_cli.cpp index 6bf3d8f..09edb58 100644 --- a/src/trellis_cli.cpp +++ b/src/trellis_cli.cpp @@ -254,6 +254,16 @@ int trellis_run(const trellis::TrellisParams& cfg) { so = trellis::shape_decode(m, slat_dn, shc, RES); m.free(); printf(" decoded voxels @res%d = %d\n", so.res, (int)so.coords.size()); mesh = trellis::dual_grid_to_mesh(so); + // feats7 is consumed by dual_grid_to_mesh and is never referenced again. + // Release it before the enormous host-side topology passes; on the + // 31M-voxel torture case this returns ~830 MiB immediately. + const size_t shape_feat_bytes = so.feats7.capacity() * sizeof(float); + std::vector().swap(so.feats7); + if (shape_feat_bytes >= (size_t)128 * 1024 * 1024) { + printf(" [host-mem] released shape feats7 %.2f GiB before mesh postprocess\n", + shape_feat_bytes / (1024.0*1024.0*1024.0)); + fflush(stdout); + } } printf(" mesh V=%d F=%d\n", mesh.V(), mesh.F()); { // reference postprocess fills small holes BEFORE the remesh (max_hole_perimeter=3e-2): @@ -289,7 +299,15 @@ int trellis_run(const trellis::TrellisParams& cfg) { trellis::Model m = trellis::Model::load(M + "/shape_dec.gguf", gpu); so_tex = trellis::shape_decode(m, lr_dn, coords, 512); m.free(); pbr_coords = &so_tex.coords; pbr_res = so_tex.res; + // The texture decoder needs only coords+subdivision masks from this + // guide; its 7-channel shape features are dead weight from here on. + std::vector().swap(so_tex.feats7); + // Mixed mode now samples PBR from so_tex.coords, so the giant HR + // coordinate table can also be returned to the OS before tex flow. + std::vector>().swap(so.coords); printf(" res-512 tex-guide decode: %d voxels\n", (int)so_tex.coords.size()); + printf(" [host-mem] mixed texture: released HR shape coords + tex-guide feats7\n"); + fflush(stdout); } // tex flow + decode inputs: HR path (shc/slat_norm/cond_dec/so.subs) vs res-512 mixed path // (coords/lr_norm/cond_512/so_tex.subs). The tex decoder upsamples via the guide subdivision. @@ -333,6 +351,10 @@ int trellis_run(const trellis::TrellisParams& cfg) { } printf(" PBR voxels=%d @res%d\n", Mv, pbr_res); } + // Texture decode has consumed the subdivision guides. Release their + // backing storage before the raw 100M+ face mesh enters weld/hole/BVH. + std::vector>().swap(so.subs); + std::vector>().swap(so_tex.subs); // `colors` is per-VOXEL but consumed per-VERTEX (weld, vertex-color GLB, PLY), // relying on dual_grid_to_mesh's vertex==voxel correspondence. fill_holes adds // cap vertices beyond Mv -- pad them (neutral grey; the caps are sub-voxel and diff --git a/src/tri_bvh.cpp b/src/tri_bvh.cpp index 8b1eeb6..0dfea7c 100644 --- a/src/tri_bvh.cpp +++ b/src/tri_bvh.cpp @@ -2,6 +2,8 @@ #include #include #include +#include +#include namespace trellis { @@ -46,6 +48,22 @@ void closest_on_tri(const float* p, const float* a, const float* b, const float* for (int k = 0; k < 3; ++k) out[k] = a[k] + ab[k]*v + ac[k]*w; } +size_t bvh_node_count_memo(int64_t n, std::unordered_map& memo) { + if (n <= 4) return 1; + auto it = memo.find(n); + if (it != memo.end()) return it->second; + const int64_t a = n / 2, b = n - a; + const size_t count = 1 + bvh_node_count_memo(a, memo) + bvh_node_count_memo(b, memo); + memo.emplace(n, count); + return count; +} + +size_t bvh_node_count(int64_t n) { + std::unordered_map memo; + memo.reserve(64); + return bvh_node_count_memo(n, memo); +} + inline float box_dist2(const float* p, const float* bmin, const float* bmax) { float d2 = 0.f; for (int k = 0; k < 3; ++k) { @@ -70,7 +88,19 @@ TriBvh TriBvh::build(const float* verts, int64_t V, const int32_t* faces, int64_ for (int k = 0; k < 3; ++k) cent[3*f+k] = (verts[3*faces[3*f]+k] + verts[3*faces[3*f+1]+k] + verts[3*faces[3*f+2]+k]) / 3.f; } - t.nodes_.reserve((size_t)F * 2); + // The old 2*F reserve is a severe over-allocation on giant decoded meshes. + // With leaf size 4 the balanced tree needs far fewer nodes; compute the exact + // count so vector growth never creates a second multi-GiB copy either. + const size_t node_need = bvh_node_count(F); + t.nodes_.reserve(node_need); + if (F >= 1000000) { + const double work_gib = ((double)node_need * sizeof(Node) + + (double)F * sizeof(int32_t) + + (double)F * 3.0 * sizeof(float)) / (1024.0*1024.0*1024.0); + printf(" [bvh-mem] F=%lld nodes=%zu build-working~%.2f GiB\n", + (long long)F, node_need, work_gib); + fflush(stdout); + } struct Span { int32_t node, begin, end; }; std::vector stack; diff --git a/src/uv_bake.cpp b/src/uv_bake.cpp index a06749a..499667b 100644 --- a/src/uv_bake.cpp +++ b/src/uv_bake.cpp @@ -507,7 +507,7 @@ int fill_holes(std::vector& verts, std::vector& faces, float max if (F == 0) return 0; // count undirected edge uses; remember one directed representative std::unordered_map> euse; // key -> {count, directed (u<<32|v)} - euse.reserve(F * 3); + euse.reserve(F * 2); auto ekey = [](int a, int b) -> uint64_t { if (a > b) { int t = a; a = b; b = t; } return ((uint64_t)(uint32_t)a << 32) | (uint32_t)b; }; for (size_t f = 0; f < F; ++f) { const int32_t* t = &faces[3*f]; @@ -611,29 +611,62 @@ void taubin_smooth(std::vector& verts, const std::vector& faces, } int fill_small_holes(std::vector& faces, int max_loop) { - const size_t F = faces.size() / 3; - auto ekey = [](int a, int b){ return ((uint64_t)(uint32_t)a << 32) | (uint32_t)b; }; - std::unordered_map dir; - dir.reserve(F * 3 * 2); - for (size_t f = 0; f < F; ++f) - for (int j = 0; j < 3; ++j) - dir[ekey(faces[3*f+j], faces[3*f+(j+1)%3])]++; - // Boundary edges traversed opposite to face winding so fan fills keep - // orientation consistent with their neighbors. Chains pass only through - // unambiguous boundary vertices (out- and in-degree exactly 1): at - // non-manifold junctions a single-successor map silently cross-links - // fragments of different holes into bogus mesh-spanning "loops". + // Memory-scaled equivalent of the old two-hash-table implementation. + // On dense 1024 meshes F can exceed 100M; keeping directed and undirected + // edge maps alive together can require tens of GiB. Both passes only need + // the boundary edge set, so collect it with ONE undirected use-count map, + // release that map, run the directed loop pass, then recollect after any + // new fan faces before the winding-agnostic pass. Boundary semantics are + // unchanged: an edge is a boundary iff its undirected use count is exactly 1. + auto dkey = [](int a, int b){ return ((uint64_t)(uint32_t)a << 32) | (uint32_t)b; }; + auto ukey = [](int a, int b){ + if (a > b) std::swap(a, b); + return ((uint64_t)(uint32_t)a << 32) | (uint32_t)b; + }; + + auto collect_boundary = [&](const char* pass) { + const size_t F = faces.size() / 3; + struct Use { int count = 0; uint64_t directed = 0; }; + std::unordered_map euse; + // This is the same asymptotic map as fill_holes(), which already runs + // successfully on the decoded mesh before texture generation. Do not + // over-reserve 6*F buckets like the previous fill_small_holes path. + euse.reserve(F * 3); + for (size_t f = 0; f < F; ++f) { + for (int j = 0; j < 3; ++j) { + const int a = faces[3*f+j], b = faces[3*f+(j+1)%3]; + auto& e = euse[ukey(a, b)]; + ++e.count; + if (e.count == 1) e.directed = dkey(a, b); + } + } + std::vector boundary; + // Boundary edges are normally tiny compared with all triangle edges; + // avoid a giant speculative reserve and let this vector grow naturally. + for (const auto& kv : euse) + if (kv.second.count == 1) boundary.push_back(kv.second.directed); + if (F >= 1000000) { + printf(" [fill-holes-stream] %s F=%zu boundary=%zu (single edge map)\n", + pass, F, boundary.size()); + fflush(stdout); + } + // Force release before allocating the small successor/adjacency maps. + decltype(euse)().swap(euse); + return boundary; + }; + + // Pass 1: directed boundary walk, same winding-preserving behavior as before. + std::vector boundary = collect_boundary("directed"); std::unordered_map nxt, outd, ind; - for (const auto& [k, cnt] : dir) { + for (uint64_t k : boundary) { const int a = (int)(k >> 32), b = (int)(uint32_t)k; - if (cnt == 1 && dir.find(ekey(b, a)) == dir.end()) { - nxt[b] = a; outd[b]++; ind[a]++; - } + nxt[b] = a; outd[b]++; ind[a]++; } std::unordered_map used; int filled = 0; size_t added = 0; for (const auto& [start, first] : nxt) { + (void)first; if (used[start] || outd[start] != 1 || ind[start] != 1) continue; std::vector loop = {start}; int cur = start; @@ -654,47 +687,49 @@ int fill_small_holes(std::vector& faces, int max_loop) { } ++filled; } - // Second pass, winding-agnostic (the GLB material is double-sided): loops - // whose boundary direction flips (simplification tears) never chain in the - // directed walk above. Assemble them over undirected boundary adjacency, - // restricted to unambiguous degree-2 vertices. - { - const size_t F2 = faces.size() / 3; - std::unordered_map und; - und.reserve(F2 * 3 * 2); - for (size_t f = 0; f < F2; ++f) - for (int j = 0; j < 3; ++j) { - const int a = faces[3*f+j], b = faces[3*f+(j+1)%3]; - und[ekey(std::min(a,b), std::max(a,b))]++; - } - std::unordered_map> adj; - for (const auto& [k, cnt] : und) { - if (cnt != 1) continue; - const int a = (int)(k >> 32), b = (int)(uint32_t)k; - adj[a].push_back(b); adj[b].push_back(a); + + // Release all pass-1 hash tables before the second full edge scan. + decltype(nxt)().swap(nxt); + decltype(outd)().swap(outd); + decltype(ind)().swap(ind); + decltype(used)().swap(used); + + // Pass 2: if pass 1 changed topology, recompute the boundary AFTER its new + // fan faces. If it filled nothing, the boundary set is unchanged and we + // can reuse it, avoiding a second 100M+-face scan entirely. + if (filled > 0) { + std::vector().swap(boundary); + boundary = collect_boundary("undirected"); + } else if (faces.size() / 3 >= 1000000) { + printf(" [fill-holes-stream] undirected: reusing boundary (directed pass filled 0)\n"); + fflush(stdout); + } + std::unordered_map> adj; + for (uint64_t k : boundary) { + const int a = (int)(k >> 32), b = (int)(uint32_t)k; + adj[a].push_back(b); adj[b].push_back(a); + } + std::unordered_map used2; + for (const auto& [start, nbrs] : adj) { + if (used2[start] || nbrs.size() != 2) continue; + std::vector loop = {start}; + int prev = start, cur = nbrs[0]; + bool cycle = false, clean = true; + for (int steps = 0; steps <= max_loop; ++steps) { + auto it = adj.find(cur); + if (it == adj.end() || it->second.size() != 2 || used2[cur]) { clean = false; break; } + if (cur == start) { cycle = true; break; } + loop.push_back(cur); + const int nx = it->second[0] == prev ? it->second[1] : it->second[0]; + prev = cur; cur = nx; } - std::unordered_map used2; - for (const auto& [start, nbrs] : adj) { - if (used2[start] || nbrs.size() != 2) continue; - std::vector loop = {start}; - int prev = start, cur = nbrs[0]; - bool cycle = false, clean = true; - for (int steps = 0; steps <= max_loop; ++steps) { - auto it = adj.find(cur); - if (it == adj.end() || it->second.size() != 2 || used2[cur]) { clean = false; break; } - if (cur == start) { cycle = true; break; } - loop.push_back(cur); - const int nx = it->second[0] == prev ? it->second[1] : it->second[0]; - prev = cur; cur = nx; - } - for (int v : loop) used2[v] = true; - if (!clean || !cycle || loop.size() < 3 || (int)loop.size() > max_loop) continue; - for (size_t i = 1; i + 1 < loop.size(); ++i) { - faces.push_back(loop[0]); faces.push_back(loop[i]); faces.push_back(loop[i+1]); - added += 1; - } - ++filled; + for (int v : loop) used2[v] = true; + if (!clean || !cycle || loop.size() < 3 || (int)loop.size() > max_loop) continue; + for (size_t i = 1; i + 1 < loop.size(); ++i) { + faces.push_back(loop[0]); faces.push_back(loop[i]); faces.push_back(loop[i+1]); + added += 1; } + ++filled; } if (filled) { printf(" fill_holes: %d boundary loops filled (+%zu faces)\n", filled, added); fflush(stdout); } return filled;