[Relax][Vision] Add get_valid_counts and classic NMS

[Aharrypotter]

Summary

Add relax.vision.get_valid_counts and classic relax.vision.non_max_suppression for object-detection post-processing pipelines.

get_valid_counts performs score-based bounding box filtering and compacts valid boxes to the front of each batch. Classic non_max_suppression performs flexible IoU-based suppression on filtered boxes, complementing existing all_class_non_max_suppression for custom post-processing workflows.

This PR implements the Relax-level registration, legalization, TOPI compute, and test coverage for both operators tracked in #18928.

Changes

Relax op registration and legalization:

• C++ op functions, FFI registration, and struct info inference for both operators (vision.h, vision.cc)
• Python wrappers with Relax docstrings (vision.py)
• Legalization to topi.vision.get_valid_counts and topi.vision.non_max_suppression
• Additional struct-info validation for score_index, id_index, and coord_start when elem_length is statically known

TOPI and testing:

• Full TOPI implementation for get_valid_counts
• Reimplementation of classic non_max_suppression in TOPI
• NumPy reference implementations in tvm.topi.testing for both operators
• Op-level tests for struct info inference, legalization, invalid attribute ranges, and e2e numerical correctness
• Stronger legalization tests that verify both relax.call_tir introduction and removal of the original Relax vision op

Limitations

• Attribute range validation for score_index, id_index, and coord_start is only enforced when the input elem_length is statically known during struct-info inference.
• Classic non_max_suppression follows the existing Relax / TOPI API shape and is intended for single-class or class-aware custom post-processing flows, distinct from all_class_non_max_suppression.

Validation

pytest tests/python/relax/test_op_vision.py -k "get_valid_counts" -v
pytest tests/python/relax/test_op_vision.py -k "test_nms_" -v

All related tests passed.

[gemini-code-assist]

Code Review

This pull request introduces the get_valid_counts and non_max_suppression operators to the Relax vision dialect. The implementation includes C++ operator registration with structural information inference, Python FFI bindings, and legalization passes to TOPI. Additionally, it provides TIR-based implementations for both operators in TOPI, along with corresponding NumPy reference implementations and comprehensive unit tests covering shape inference, legalization, and end-to-end execution. I have no feedback to provide.

[tlopex]

Still some points to be improved

• The iou_threshold type check only handles float, so iou_threshold=0 will not be converted to a TIR constant. Please use isinstance(iou_threshold, (float, int)), consistent with get_valid_counts.
• When id_index < 0, there are no class IDs in the input, so force_suppress=False cannot actually enforce same-class-only suppression. In practice it ends up behaving the same as force_suppress=True. This is probably fine for compatibility, but it would be helpful to call it out in the docstring.
• Test coverage still seems incomplete: no case for force_suppress=True, no case where max_output_size > 0 actually truncates outputs, no multi-batch NMS test, and no coverage for the return_indices=False, invalid_to_bottom=False path.

[Aharrypotter]

PTAL @tlopex

[tlopex]

A few follow-ups:

• Please update _rearrange_out to use T.parallel(0, batch_size) as well;
• Please tighten id_index validation to reject values less than -1, e.g. if (attrs->id_index < -1 || attrs->id_index >= elem_length), since -2, -3, etc. currently pass silently.
• The score_threshold type handling works, but the current double isinstance check reads a bit awkwardly; please consider simplifying it.
• Please consider adding a test for invalid_to_bottom=True, return_indices=True, even if invalid_to_bottom is a no-op on that path.

You can have a look and resolve the existed conflicts

[Aharrypotter]

Thanks for the detailed follow-ups, @tlopex ! I'm currently digesting these suggestions and working on the fixes.

It might take me a little bit of time to get everything updated, tested properly, and resolve the conflicts. I'll ping you once the new commits are ready.

[Aharrypotter]

PTAL @tlopex

[tlopex]

LGTM Thanks!