Support NMS op lowering

[jinchen62]

No description provided.

[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.