✨[Feature] Do not remove nodes which are in-place operations (classified in TRT segments) in fallback

[peri044]

Is your feature request related to a problem? Please describe. Currently for eg: if aten::append is in prim::If block, it gets removed because it is classified as TRT segment. No output is registered for it because aten::append is an inplace op. Make modifications to registerSegmentsOutputs

Segmented Graph: graph(%self_1 : __torch__.ModifyingOpError_trt,
      %x.1 : Tensor,
      %y.1 : Tensor):
  %mod_list.1 : Tensor[] = prim::ListConstruct(%x.1)
  %__torch___ModifyingOpError_trt_engine_0x557361771710 : __torch__.torch.classes.tensorrt.Engine = prim::GetAttr[name="__torch___ModifyingOpError_trt_engine_0x557361771710"](%self_1)
  %5 : Tensor[] = prim::ListConstruct(%x.1, %y.1)
  %6 : Tensor[] = tensorrt::execute_engine(%5, %__torch___ModifyingOpError_trt_engine_0x557361771710)
  %7 : Tensor = prim::ListUnpack(%6)
  %8 : bool = aten::Bool(%7) # mod.py:12:11
   = prim::If(%8)
    block0():
      -> ()
    block1():
      -> ()
  %9 : int = prim::Constant[value=0]()
  %10 : Tensor = aten::cat(%mod_list.1, %9) # mod.py:14:12
  return (%10)

Describe the solution you'd like

Describe alternatives you've considered

Additional context

[peri044]

Also add this as a test case for this PR https://github.com/pytorch/TensorRT/pull/1067

class ModifyingOpError(nn.Module):
    def __init__(self):
        super(ModifyingOpError, self).__init__()
    def forward(self, x, y):
        mod_list = [x]
        if x.sum() > y.sum():
            mod_list.append(y)
        z = torch.cat(mod_list)
        return z

[ncomly-nvidia]

Resolved by #1140. Closing

[bowang007]

@ncomly-nvidia I don't this this one should be closed. It refers to whether we should delete these several lines https://github.com/pytorch/TensorRT/blob/0738caacf8d0911521ff73b0f67b071b8c2b7b42/core/partitioning/partitioning.cpp#L202

[github-actions[bot]]

This issue has not seen activity for 90 days, Remove stale label or comment or this will be closed in 10 days

[narendasan]

@gs-olive / @bowang007 How does this relate to index_put_?

[gs-olive]

@gs-olive / @bowang007 How does this relate to index_put_?

I don't think this specific case is directly related to index_put_ other than that both append and index_put_ are in-place operations. In this case, the issue seems to be that no-return if...else code blocks are removed, while they could still be making meaningful changes to existing tensors (in-place). Regarding index_put_, if it were located within an if...else block without a return value in TorchScript like the above, it would also be optimized away, though this holds true for any in-place operation.

I think in both the case of append and index_put_, functionalizing the graph as much as possible is very important, especially considering proposed new features including allowing Int64 inputs (RFC #1553), or debugging in-place issues related to assigning values to TRT Constant-labeled tensors, as in Issue #1453

[github-actions[bot]]

This issue has not seen activity for 90 days, Remove stale label or comment or this will be closed in 10 days