[TOPI] Improve memory layout inside GPU NMS kernel

[masahi]

Motivation

Currently, our NMS API expects input and output data to be in a format [batch, num_anchors, 6], where 6 is for 4 bbox coordinates, score and class id. This format is highly inefficient for class-aware NMS, because to check if two bboxes are the same class, we need to load other 5 values. More discussion in #7154 (comment) and other comments in that PR.

This PR improved the memory layout of buffers used in the NMS kernel. Specifically, it unpacks [batch, num_anchors, 6] input into one [batch, num_anchors, 4] buffer for bboxes, and two [batch, num_anchors] buffers for scores and class ids.

Speed up

This change is expected to give a good speed up for the cases when the number valid boxes is large and there are class ids involved. For example, although the change in #7154 brings only modest speed up, combined with this PR the speed up becomes good, as shown below:

PyTorch NMS workload from FasterRCNN / MaskRCNN (4500 boxes):

Current main : 5.78 milli sec
After rewrite in #7154 to restore class aware NMS: 5.23 milli sec
With memory layout improvement in this PR and the rewrite in #7154: 4.51 milli sec

Other consideration

For other cases, it is likely that there would be no speed up, but it should be no worse. For example, it doesn't improve NMS workload from Gluon SSD. When return_indices=False, which I think only applies to MXNet, there is additional concatenation of three buffers at the end. For Gluon SSD, this concat takes 40 micro sec on GTX 1070 ti, while the main NMS kernel takes 200 micro sec. If people think this additional concat is expensive, we really need to rethink our NMS API: The same concat is already happenning in all other frontends, because all frameworks other than MXNet gives NMS inputs in an unpacked form and concat is necessary to satisfy our NMS API.

please review @mbrookhart @Laurawly @vinx13 @trevor-m

[mbrookhart]

1. I'm glad this speeds up NMS, that confirms some of our suspicions from [Torch] Restore class-aware NMS for detection models by graph rewrite #7154
2. I don't love that after this PR, ever framework but MXNET will have an op of the form: concat->split->nms->concat->split. Can we talk about moving the split and concat steps out of this op and reworking the frameworks for a new API?

[mbrookhart]

Even after this, I think we will still need a loop over classes for ONNX and TF, since ONNX explicitly and TF implicitly need max_output_boxes_per_class, while this op even with class id will return max_output_boxes for all classes.

[masahi]

I don't quite follow, maybe you are missing something?

First, this PR doesn't change our NMS API, it only changes the buffer layouts used internally.

Second, the final concat is only required for MXNet, which uses return_indices=False. Our NMS returns something completely different depending on return_indices flag : If True, it returns a big output tensor packed with bboxes, scores and class ids, with invalid entries indicated by -1.

(The valid entries are supposed to move to the top, if invalid_to_bottom flag is True. But our GPU NMS kernel ignores this argument and the output is not reordered. This is another difference with CPU implementation, I think this is a bug)

tvm/python/tvm/topi/cuda/nms.py

Lines 762 to 763 in 4e8cc4f

	invalid_to_bottom : optional, boolean
	Whether to move all valid bounding boxes to the top.

If return_indices=True, which applies to TF, ONNX, and PyTorch, we only return survived box indices, so there is no need to concat bboxes, scores, and class ids. For this case, there is zero additional overhead after this PR, it is just that concat before NMS done by frontends are now completely pointless.

[mbrookhart]

I see, yep.

We don't seem to have a test for invalid_to_bottom, it's False by default:

tvm/python/tvm/relay/op/vision/nms.py

Line 70 in 4e8cc4f

invalid_to_bottom=False,

and the only test that mentions it sets it to False (in test_any). Grep-ing around the codebase, the only place it's ever set as True is in the MXNet frontend:

tvm/python/tvm/relay/frontend/mxnet.py

Line 1050 in 4e8cc4f

invalid_to_bottom=True,

and I don't see a unit test for that op.

It might be a bug because it never runs outside of the gluon ssd tutorial, which doesn't check box order:

tvm/tutorials/frontend/deploy_ssd_gluoncv.py

Lines 111 to 127 in 926a315

	for target in ["llvm", "cuda"]:
	ctx = tvm.context(target, 0)
	if ctx.exist:
	lib = build(target)
	class_IDs, scores, bounding_boxs = run(lib, ctx)
	######################################################################
	# Display result
	ax = utils.viz.plot_bbox(
	img,
	bounding_boxs.asnumpy()[0],
	scores.asnumpy()[0],
	class_IDs.asnumpy()[0],
	class_names=block.classes,
	)
	plt.show()

We should probably open an issue for it, or add a unit test and fix the kernel.

I'd still like to see us change the NMS API to better reflect what's going on internally (i.e., remove the concat/split for TF/ONNX/Pytorch), but that doesn't need to be this PR.

[masahi]

Yes, ideally I want to update our NMS to be closer to TF/ONNX/PyTorch, and let MXNet frontend handle split and concat, rather than the other way around (what we have now). Current API is over complicated due to the need to support both styles. If we can assume that return_indices is always True, we can clean up our API a lot. For example, invalid_to_bottom argument only makes sense for MXNet. We don't need coord_start, score_index, and id_index arguments, if inputs are only unpacked.

Supporting max_output_boxes_per_class needs change in the implementation as well. We need to count the number of survived boxes per class. But that's the only change I think, it is definitely doable without writing another kernel.

[masahi]

Thanks @mbrookhart