Refactor nms into TorchVision variant.

[ysiraichi]

This PR refactors the existing nms lowering implementation, changing the appropriate parts for complying with TorchVision semantics. In summary:

• Implement new nms lowering, based on the old implementation
  • More comments
  • TorchVision semantics
• Adapts TorchVision tests for nms
• Move the lowering implementation to torch_xla/csrc/xla_lower_util.cpp
  • Alongside other lowering implementations
• Register the new kernel as an XLA dispatch for torchvision::nms
  • Make it possible to call the kernel by calling torchvision.ops.nms(...) directly
• Kills the old implementation

cc @miladm @JackCaoG

[ysiraichi]

I'm having a bit of trouble identifying the cause of the CI failure. Basically, if we set XLA_USE_EAGER_DEBUG_MODE=1, the test TestNMS::test_nms_ref starts failing with the following error:

RuntimeError: Bad StatusOr access: INVALID_ARGUMENT: Executable expected parameter 0 of size 4000
but got buffer with incompatible size 4004

After looking into where this error was coming from, I found out that it was due to PyTorch/XLA trying to figure out the actual size of the output.

My guess is that, since I call xla::SetDimensionSize in the lowering, one of the dimensions are marked as dynamic. Which makes the XLATensorImpl::sym_sizes_ be populated with a c10::SymInt. When we try to retrieve the actual size (triggered by a to_functional_tensor(t) call on the result of nms), PyTorch/XLA tries to figure it out by actually running the computation.

---

Given all this information, I'm not sure where the "buffer with imcompatible size 4004" is coming from. test_nms_ref does create a tensor of size 4000 (1000x4 elements). But, the last 4 is unclear to me.

Another thing I'm puzzled about is why is this test failing only when XLA_USE_EAGER_DEBUG_MODE is set? Why wouldn't a normal test run trigger the same error?

@JackCaoG Have you ever seen an error like this? I believe other lowerings also call this same function (e.g. BuildMaskedSelect). So, doesn't seem like a problem with xla::SetDimensionSize being used.

[JackCaoG]

XLA_USE_EAGER_DEBUG_MODE it is a debug thing where it tries to execute the value of the node upon construction. You can think of it as calling mark_step every time a IR node is created.

[ysiraichi]

Right. So, what is puzzling is why would it fail only when XLA_USE_EAGER_DEBUG_MODE is set. Do you have any guesses?

[JackCaoG]

not really but feel free to skip the test when XLA_USE_EAGER_DEBUG_MODE is enabled. nms is a rarely used ops that's not even aten, XLA_USE_EAGER_DEBUG_MODE is a debug only flag. I think it is OK.

[vanbasten23]

I wonder if the implementation here is a copy with some modification of the deleted nms_op.cpp, or if it is a brand new implementation.

[ysiraichi]

I would say that it is a new implementation heavily based on the previous one. So, the problems that I saw with the previous implementation were:

• Almost no comments: couldn't easily tell what some parts of the code were doing
• Copied from another source: it was copied from an old version of tensorflow
• Different signature and semantics: it returned the best output_size box indices, even though some of them might have been suppressed

I believe this implementation checks all of the above-mentioned boxes. I added plenty of comments + a few changes that resulted in a more maintainable code (IMHO).

[ysiraichi]

@JackCaoG @vanbasten23 I think this PR is ready. Could you take a look at it when you have some time?

[vanbasten23]

Is there any resource that you referenced when coming up with the implementation?

[ysiraichi]

I based this implementation in two other implementations:

• The old PyTorch/XLA nms implementation
• TorchVision CUDA nms implementation

[miladm]

@ysiraichi can you please add a dynamic shape input test scenario for nms? Because this op is dynamic, it's been falling back to CPU. Of course, there is a lot of interest to bring the op to the xla device; though, this requires correct functionality for dynamism.

[miladm]

more clarity: number of boxes is set to 1000 at the moment. we want that number to be dynamic.

[miladm]

to be clear the output of this test appears to be dynamic, though the input number of boxes can also be dynamic for nms. This test, currently, covers one dynamism scenario (i.e. the output dynamism).