Support NMS op lowering

[jinchen62]

TODO: support multiple batches and classes

[zjgarvey]

Like we were talking about on a call, I definitely think we should tensorize some of this.

In the outer loop, we pick a single box to compare to the other boxes.
Instead of computing the IOU for each other box individually in the inner loop, lets do:

// outside all loops:
Value x1Slice /*=slice boxes to get x1 values for each box*/;
Value x2Slice ..;
Value y1Slice ..;
Value y2Slice ..;
Value xDistance = x2Slice - x1Slice;
Value yDistance = y2Slice - y1Slice;
Value boxAreas = xDistance * yDistance;
// inside loop over sorted boxes:
Value currBox /*=the box with the highest score among available boxes*/;
// The elementwise tensor arithmetic below allows broadcasting
Value innerX1 = max(x1Slice, currBox[0]);
Value innerX2 = min(x2Slice, currBox[2]);
Value innerY1 ..;
Value innerY2 ..;
Value intersectionDistanceX = innerX2 - innerX1;
Value intersectionDistanceY = innerY2 - innerY1;
Value intersectionArea = intersectionDistanceX * intersectionDistanceY;
Value currArea = boxAreas[currBoxIdx];
Value unionArea = boxAreas + currArea - intersectionArea;
Value IOU = intersectionArea / unionArea;

Actually, based on the ordering of [x1,y1,x2,y2] we could even do:

// outside all loops:
lowSlice = [x1,y1] slice;
highSlice = [x2,y2] slice;
distances = highSlice - lowSlice;
area = reduceProd(distances, the dim that has size two);
// inside loop
innerLow = max(lowSlice, currBox[x1,y1]);
innerHigh = min(highSlice, currBox[x2,y2]);
innerDistance = innerHigh - innerLow;
intersectionArea = reduceProd(innerDistance);
currArea = area[currBoxIdx];
unionArea = area + currArea - intersectionArea;
IOUs = intersectionArea / unionArea;

Although, I'm not sure how we can go about skipping the IOU calculations that are redundant. I'm actually curious if it is worth it to skip redundant computations if it requires us to extract individual elements and do the arithmetic one at a time. At the very least, computing each of the box areas outside the loops is going to be an improvement.

[jinchen62]

Three e2e tests pass with #3892 and without flag --iree-input-demote-f64-to-f32=false(iree-org/iree#19299).

• test_nonmaxsuppression_identical_boxes
• test_nonmaxsuppression_single_box
• test_nonmaxsuppression_suppress_by_IOU

[zjgarvey]

The implementation looks good to me, and if we can pass e2e tests, I think this is good.

My only lingering concern is related to the signature of the onnx op when one of the other optional args is missing. It might be good to add some test cases to check what happens there and make sure we won't mis-compile those examples.

Another point of concern is that this whole implementation could likely be done in the onnx-to-torch lowering and might allow support for the batched version of NMS that onnx uses (we have some unsqueezes in this conversion that cannot reasonably be assumed to succeed). I don't think we need to make that change now, but it's something to keep in mind if we end up seeing this op in models we care about with non-trivial batching.