Add vector-Jacobian-products for a subset of nvFuser-supported prims; add backward support for nvprims
Labels: triaged, open source, cla signed, module: nvfuser, module: primTorch, no-stale

NOTE: Currently, this PR is blocked until https://github.com/pytorch/pytorch/issues/85696 is resolved.

`vjp_implementations["prim_name"]` is a callable that computes vector-Jacobian products for given cotangents, forward results, and forward inputs. And this function is used in the `torch.autograd.Function.backward` method.

Here's an example trace for the `torch.sigmoid` function:
```py
import torch
from torch._prims.context import NvfuserPrimsMode, TorchRefsMode
from torch.fx.experimental.proxy_tensor import make_fx

a = torch.randn(3, 3, device='cuda', requires_grad=True)
g = torch.ones_like(a)

def func(a):
    c = torch.sigmoid(a)
    return torch.autograd.grad(c, a, g)

# torch._refs.sigmoid trace
with NvfuserPrimsMode(), TorchRefsMode():
    make_fx(func)(a).graph.print_tabular()
```
```
opcode         name                 target                     args                        kwargs
-------------  -------------------  -------------------------  --------------------------  --------
placeholder    a_1                  a_1                        ()                          {}
call_function  neg                  nvprims.neg.default        (a_1,)                      {}
call_function  detach               aten.detach.default        (neg,)                      {}
call_function  exp                  nvprims.exp.default        (neg,)                      {}
call_function  detach_1             aten.detach.default        (exp,)                      {}
call_function  add                  nvprims.add.default        (1.0, exp)                  {}
call_function  detach_2             aten.detach.default        (add,)                      {}
call_function  div                  nvprims.div.default        (1.0, add)                  {}
call_function  detach_3             aten.detach.default        (div,)                      {}
get_attr       _tensor_constant0    _tensor_constant0          ()                          {}
call_function  is_same_size         aten.is_same_size.default  (div, _tensor_constant0)    {}
call_function  detach_4             aten.detach.default        (detach_3,)                 {}
call_function  detach_5             aten.detach.default        (detach_3,)                 {}
get_attr       _tensor_constant0_1  _tensor_constant0          ()                          {}
call_function  div_1                nvprims.div.default        (_tensor_constant0_1, add)  {}
get_attr       _tensor_constant0_2  _tensor_constant0          ()                          {}
call_function  neg_1                nvprims.neg.default        (_tensor_constant0_2,)      {}
call_function  mul                  nvprims.mul.default        (neg_1, 1.0)                {}
call_function  pow_1                nvprims.pow.default        (add, -2)                   {}
call_function  mul_1                nvprims.mul.default        (mul, pow_1)                {}
call_function  detach_6             aten.detach.default        (detach_2,)                 {}
call_function  detach_7             aten.detach.default        (detach_2,)                 {}
call_function  view_of              nvprims.view_of.default    (mul_1,)                    {}
call_function  view_of_1            nvprims.view_of.default    (mul_1,)                    {}
call_function  detach_8             aten.detach.default        (detach_1,)                 {}
call_function  detach_9             aten.detach.default        (detach_1,)                 {}
call_function  mul_2                nvprims.mul.default        (view_of_1, detach_9)       {}
call_function  detach_10            aten.detach.default        (detach,)                   {}
call_function  detach_11            aten.detach.default        (detach,)                   {}
call_function  neg_2                nvprims.neg.default        (mul_2,)                    {}
output         output               output                     ((neg_2,),)                 {}
```

`aten.detach` and `aten.is_same_size` appear in the trace that are called somewhere in the autograd engine code but other than that the trace consists of `nvprims` calls.

Thanks to @rdspring1 for implementing several VJP functions!


cc @kevinstephano @jjsjann123 @ezyang @mruberry @ngimel @Lezcano @fdrocha @peterbell10