Implement setitem without copy_

[shino16]

Fixes #2790. Previously, inp[idx] = value was translated into _copy_(inp, clang.copy_with_setitem(inp, idx, val)), where copy_with_setitem was an opaque primitive that only torchex could execute. This PR makes the entire setitem opaque.

Demo

import torch, thunder

def fn(x, idx, val):
    x[idx] = val

jf = thunder.jit(fn)
x = torch.zeros(3, 4, 5, device="cpu")
x_ref = x.clone()
idx = torch.tensor([1], device="cpu")
val = torch.ones(1, 4, 5, device="cpu")
fn(x_ref, idx, val)
jf(x, idx, val)

print(thunder.last_traces(jf)[-1])
torch.testing.assert_close(x, x_ref)

# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x, idx, val):
  # x: "cpu f32[3, 4, 5]"
  # idx: "cpu i64[1]"
  # val: "cpu f32[1, 4, 5]"
  (t4,) = update_aliases((x,))
  del x

  # /opt/pytorch/lightning-thunder/tmp/main.py:4:           x[idx] = val
  t5 = setitem(t4, idx, val)  # t5: "cpu f32[3, 4, 5]"
    # t5 = ltorch.setitem_(t4, idx, val)  # t5: "cpu f32[3, 4, 5]"
      # t5 = prims.setitem(t4, idx, val)  # t5: "cpu f32[3, 4, 5]"
  del t4
  return ()

[kiya00]

The failures in CI seems relevant

[shino16]

We have to define forward and backward separately. If not, when we extract the forward trace from the combined fw&bw function, we hit into prims.setitem(clone(g), index, 0) in backward, which survives through DCE because setitem is tagged as DONT_DCE.