How to remove in-place copy_ operator from the traced graph?

I traced one of the detectron2 models and noticed that the resulting graph has in-place copy_ operators.

Is it possible to somehow remove/replace this copy_ operator from the graph?
Maybe jit has an option to produce a graph without in-place operators (such as copy_)?

Probably it is better to rewrite the following 4 lines in the model itself.
Detectron2 - box2box_transform.apply_deltas:111

Changing these lines of code might remove the copy_ ops and you could try it by locally creating this model with the changes.
However, just out of interest, why would you like to remove these operations from the graph?

I think the background is moving the model to TVM see this discussion on TVM and copy_.

@apivovarov You probably increase your chance of getting good answers if you give more context.

Personally, I think the two main ways of doing this is:

  • As you suggest, the immediate way is to change the model to use torch.cat + view instead. This is probably wildly preferable because any translation of index assignment look funny.
  • Implement view tracking in TVM. I commented a bit on this in the TVM copy_ PR before yours. Takes some effort but after that it should be reasonably possible to express index assignments. There will likely be corner cases that you miss as there are operations PyTorch (e.g. reshape) that may or may not return views.

Best regards

Thomas

I think probably the best way to go about this is to prove that the input is linearly typed and replace copy_ with a functional op that creates a copy of the tensor and then writes to it. This is basically how other functional tensor languages handle writes - like jax and numlin. You need to prove that the input is linearly typed (there are no other aliases to it) so that you don’t change the semantics of the program.

I maintain a pass that does this logic for many inplace ops but does not yet handle copy_. I would accept a PR that adds a new operator similar to jax.ops.index_update and replaces copy_ with it when that can be proved safe.

Looking at apply_deltas, we initialize the tensor and then write to it so we would successively be able to remove copy_ from this function.

Another place which probably should be rewritten in Detectron2 is boxes.clip() - detectron2/boxes.py at master · facebookresearch/detectron2 · GitHub