Batch pose NMS tiles and skip suppressed references

[Yuan-Xinyi]

Description

This PR accelerates grasp-pose NMS (embodichain/utils/nms.py) by ~40×, cutting OpenDoor plan compilation from 10.6 s to 0.72 s end-to-end. Semantics are frozen: public API, thresholds, visit order, tie-breaking, and output index sequences are unchanged.

Why NMS dominated planning

Profiling OpenDoor compilation (UR5, default microwave tutorial) showed the motion stack itself is healthy — the 10.6 s went almost entirely elsewhere:

stage	time	share
grasp-pose NMS (pose_nms)	10.27 s	90%
IK (169 calls, analytic URSolver)	0.16 s	1.5%
trajectory generation (Toppra ×4)	0.17 s	1.6%

The antipodal sampler produces ~43k grasp candidates on the door handle; NMS reduces them to a few hundred survivors.

Root cause: cubic launch scaling + single-threaded CPU tiles

The greedy pass computed pairwise tiles via one Warp launch per host-bounded reference block, with the block budget chunk²//N shrinking as N grows — total launches scale as N³/chunk³ (≈10,200 at N=43k, matching the 9,959 observed in the profile). Two additional measurements rule out a machine-specific cause:

• per-launch overhead on this machine is a normal 0.1–1 ms (CPython arg marshalling); a faster CPU halves the constant but not the cubic growth (N=90k ⇒ ~86k launches ⇒ minutes on any machine);
• the grasp pipeline hands NMS CPU tensors, so each Warp tile also executed single-threaded on the CPU.

What changed

1. Batched torch tiles replace per-tile kernel launches: identical elementwise float32 threshold math (quaternion |dot| for rotation, squared Euclidean for translation) with the same reduction order.
2. Alive filtering: references are still visited strictly in the configured order, but a reference already suppressed when its block starts can never be kept, so its closeness row is never computed. Cost now scales with the number of survivors (43k → ~900) instead of the raw candidate count.
3. CUDA offload for CPU inputs: the pairwise math runs on CUDA when available (inputs are ~1 MB); decisions are identical and indices return on the input device.

Measured results

benchmark	before	after
OpenDoor NMS stage (43,144 candidates, real handle mesh)	10.27 s	0.21–0.31 s
OpenDoor engine.compile end-to-end	10.60 s	0.72 s
synthetic sweeps 5k–44k poses, old vs new index sequences	—	exactly equal in all runs

Every skill that samples grasps through graspkit (pick_up, place, hand_over, axis_align, open_door, slide, twist, …) benefits; larger --n_sample settings benefit the most.

Equivalence enforcement

tests/utils/test_nms.py gains a literal O(N²) reference implementation of the documented semantics and asserts exact index-sequence equality across: both orderings (preserve_order and neighbor-count priority with index tie-break), chunk sizes 1/7/128/2048, the angle_th > π rotation-always-close branch, an all-duplicates input, a 6k heavy-suppression profile matching real grasp-candidate statistics (also validating the CUDA offload against the CPU reference), and CUDA-resident inputs. The 7 pre-existing behavioral tests pass unchanged, as do the 6 graspkit pg_grasp tests downstream.

Dependencies: none (removes the module's Warp dependency).

Type of change

• Enhancement (non-breaking change which improves an existing functionality)

Screenshots

N/A

Checklist

• I have run the black . command to format the code base.
• I have made corresponding changes to the documentation (no routed docs cover utils/nms.py; behavior unchanged)
• Public API changes are reflected in the API docs (no API change; python docs/scripts/check_api_docs.py: 1925/1925)
• I have added tests that prove my feature works (see equivalence enforcement above)
• Dependencies have been updated, if applicable (none required)

Validation

pytest tests/utils/test_nms.py             -> 37 passed
pytest tests/toolkits/test_pg_grasp.py     -> 6 passed (downstream consumer)
python docs/scripts/check_api_docs.py      -> 1925/1925
black .                                    -> clean

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the previously reported CUDA fallback, memory-bound, and threshold-association concerns are resolved in the current code.

Findings

1.  Tiles No Longer Bound Memory ▶

Fix with agent prompt

### Issue 1
embodichain/utils/nms.py:undefined-77
`_close_block` retains an `(R, N)` matrix where `R` can equal `chunk_size`, so peak storage grows as `chunk_size × num_poses` during both neighbor counting and suppression. Large valid sample configurations can therefore exhaust CPU or GPU memory even though `chunk_size` is documented as bounding both dimensions of a pairwise tile.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

• Preserves greedy visit order, tie-breaking, thresholds, and returned index devices.
• Bounds pairwise working storage to approximately chunk_size² entries.
• Adds reference-equivalence coverage across ordering modes, tile sizes, duplicate-heavy inputs, rotation-only behavior, and CUDA execution.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Input pose matrices] --> B[Extract float32 positions and quaternions]
    B --> C{CPU input and CUDA available?}
    C -->|Yes| D[Run bounded pairwise tiles on CUDA]
    D -->|CUDA RuntimeError| E[Warn and retry on CPU]
    C -->|No| E
    D -->|Success| F[Count neighbors when priority ordering is requested]
    E --> F
    F --> G[Visit references in configured order]
    G --> H[Skip references already suppressed]
    H --> I[Compute bounded closeness tiles]
    I --> J[Return kept original indices on input device]

Loading

_{Reviews (2) · Last reviewed commit: "fix(utils): bound NMS tiles, stabilize a..."}

[greptile-apps]

Tiles No Longer Bound Memory

_close_block retains an (R, N) matrix where R can equal chunk_size, so peak storage grows as chunk_size × num_poses during both neighbor counting and suppression. Large valid sample configurations can therefore exhaust CPU or GPU memory even though chunk_size is documented as bounding both dimensions of a pairwise tile.

Knowledge Base Used: Runtime configuration and utilities

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/utils/nms.py
Line: 77

Comment:
**Tiles No Longer Bound Memory**

`_close_block` retains an `(R, N)` matrix where `R` can equal `chunk_size`, so peak storage grows as `chunk_size × num_poses` during both neighbor counting and suppression. Large valid sample configurations can therefore exhaust CPU or GPU memory even though `chunk_size` is documented as bounding both dimensions of a pairwise tile.

**Knowledge Base Used:** [Runtime configuration and utilities](https://app.greptile.com/dexforce/-/custom-context/knowledge-base/dexforce/embodichain/-/docs/runtime-configuration-and-utilities.md)

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[Yuan-Xinyi]

Fixed in 17234c8. The greedy pass sizes reference blocks by the same chunk_size**2 // num_poses entry budget the pre-rewrite implementation used (the old slowness came from per-tile kernel launches, not block granularity — with batched torch tiles the extra outer iterations are negligible: OpenDoor end-to-end stays at 0.30 s). Neighbor counting no longer materializes any (rows, num_poses) matrix at all: counts accumulate per bounded tile, with the diagonal cleared per tile by index. chunk_size again bounds both dimensions of every pairwise allocation.