[PrimTorch] Functionalization pass removes Instance Norm / Batch Norm running stats transformations
Labels: high priority, triaged, module: primTorch

### 🐛 Describe the bug

The alias ops on running stats for Instance Norm are removed after functionalization pass. NvFuser uses the alias ops to map the running stats back to the fusion input's original shape, stride, and dtype.

References:
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Normalization.cpp#L664-L670
https://github.com/pytorch/pytorch/blob/master/functorch/_src/aot_autograd.py#L393-L413

```python
====== Input graph ======
class joint_forward_backward(torch.nn.Module):
    def forward(self, primals, tangents):
        primals_1: f32[32, 32, 1, 1], primals_2: f32[32], primals_3: f32[32], primals_4: i64[], primals_5: f32[2, 32, 8, 8], tangents_1: f32[2, 32, 8, 8], = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:20, code: out = self.conv(inp)
        convolution: f32[2, 32, 8, 8] = torch.ops.aten.convolution.default(primals_5, primals_1, None, [1, 1], [0, 0], [1, 1], False, [0, 0], 1)

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:21, code: out = out.relu()
        relu: f32[2, 32, 8, 8] = torch.ops.aten.relu.default(convolution);  convolution = None
        detach: f32[2, 32, 8, 8] = torch.ops.aten.detach.default(relu)

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:22, code: out = self.bn(out)
        repeat: f32[64] = torch.ops.aten.repeat.default(primals_2, [2])
        repeat_1: f32[64] = torch.ops.aten.repeat.default(primals_3, [2])
        view: f32[1, 64, 8, 8] = torch.ops.aten.view.default(relu, [1, 64, 8, 8]);  relu = None
        empty: u8[0] = torch.ops.aten.empty.memory_format([0], dtype = torch.uint8, layout = torch.strided, device = device(type='cuda', index=0))
        native_batch_norm = torch.ops.aten.native_batch_norm.default(view, None, None, repeat, repeat_1, True, 0.1, 1e-05)
        getitem: f32[1, 64, 8, 8] = native_batch_norm[0]
        getitem_1: f32[64] = native_batch_norm[1]
        getitem_2: f32[64] = native_batch_norm[2];  native_batch_norm = None

        # ===========================================
        # running mean for instance_norm
        alias: f32[32] = torch.ops.aten.alias.default(primals_2);  primals_2 = None
        view_1: f32[2, 32] = torch.ops.aten.view.default(repeat, [2, 32])
        mean: f32[32] = torch.ops.aten.mean.dim(view_1, [0]);  view_1 = None
        copy_: f32[32] = torch.ops.aten.copy_.default(alias, mean);  alias = mean = None
        # ===========================================

        # ===========================================
        # running var for instance_norm
        alias_1: f32[32] = torch.ops.aten.alias.default(primals_3);  primals_3 = None
        view_2: f32[2, 32] = torch.ops.aten.view.default(repeat_1, [2, 32])
        mean_1: f32[32] = torch.ops.aten.mean.dim(view_2, [0]);  view_2 = None
        copy__1: f32[32] = torch.ops.aten.copy_.default(alias_1, mean_1);  alias_1 = mean_1 = None
        # ===========================================

        view_3: f32[2, 32, 8, 8] = torch.ops.aten.view.default(getitem, [2, 32, 8, 8]);  getitem = None

        # No stacktrace found for following nodes
        is_same_size = torch.ops.aten.is_same_size.default(view_3, tangents_1)

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:22, code: out = self.bn(out)
        view_4: f32[1, 64, 8, 8] = torch.ops.aten.view.default(tangents_1, [1, 64, 8, 8]);  tangents_1 = None
        native_batch_norm_backward = torch.ops.aten.native_batch_norm_backward.default(view_4, view, None, repeat, repeat_1, getitem_1, getitem_2, True, 1e-05, [True, False, False]);  view_4 = view = repeat = repeat_1 = getitem_1 = getitem_2 = None
        getitem_3: f32[1, 64, 8, 8] = native_batch_norm_backward[0]
        getitem_4 = native_batch_norm_backward[1]
        getitem_5 = native_batch_norm_backward[2];  native_batch_norm_backward = None
        view_5: f32[2, 32, 8, 8] = torch.ops.aten.view.default(getitem_3, [2, 32, 8, 8]);  getitem_3 = None

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:21, code: out = out.relu()
        detach_1: f32[2, 32, 8, 8] = torch.ops.aten.detach.default(detach);  detach = None
        threshold_backward: f32[2, 32, 8, 8] = torch.ops.aten.threshold_backward.default(view_5, detach_1, 0);  view_5 = detach_1 = None

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:20, code: out = self.conv(inp)
        convolution_backward = torch.ops.aten.convolution_backward.default(threshold_backward, primals_5, primals_1, [0], [1, 1], [0, 0], [1, 1], False, [0, 0], 1, [False, True, False]);  threshold_backward = primals_5 = primals_1 = None
        getitem_6 = convolution_backward[0]
        getitem_7: f32[32, 32, 1, 1] = convolution_backward[1]
        getitem_8 = convolution_backward[2];  convolution_backward = None
        return pytree.tree_unflatten([view_3, getitem_7, None, None, None, None], self._out_spec)
```

Joint graph after functionalization pass:
https://github.com/pytorch/pytorch/blob/master/functorch/_src/aot_autograd.py#L419
```python
====== Joint graph ======
class fake_fn(torch.nn.Module):
    def forward(self, primals, tangents):
        primals_1: f32[32, 32, 1, 1], primals_2: f32[32], primals_3: f32[32], primals_4: i64[], primals_5: f32[2, 32, 8, 8], tangents_1: f32[2, 32, 8, 8], = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:20, code: out = self.conv(inp)
        convolution: f32[2, 32, 8, 8] = torch.ops.aten.convolution.default(primals_5, primals_1, None, [1, 1], [0, 0], [1, 1], False, [0, 0], 1)

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:21, code: out = out.relu()
        relu: f32[2, 32, 8, 8] = torch.ops.aten.relu.default(convolution);  convolution = None
        detach_1: f32[2, 32, 8, 8] = torch.ops.aten.detach.default(relu)
        detach_2: f32[2, 32, 8, 8] = torch.ops.aten.detach.default(detach_1);  detach_1 = None

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:22, code: out = self.bn(out)
        repeat: f32[64] = torch.ops.aten.repeat.default(primals_2, [2]);  primals_2 = None
        repeat_1: f32[64] = torch.ops.aten.repeat.default(primals_3, [2]);  primals_3 = None
        view: f32[1, 64, 8, 8] = torch.ops.aten.view.default(relu, [1, 64, 8, 8]);  relu = None
        native_batch_norm = torch.ops.aten.native_batch_norm.default(view, None, None, repeat, repeat_1, True, 0.1, 1e-05)
        getitem: f32[1, 64, 8, 8] = native_batch_norm[0]
        getitem_1: f32[64] = native_batch_norm[1]
        getitem_2: f32[64] = native_batch_norm[2];  native_batch_norm = None
        
        # ===========================================
        # view_1 and view_2 are dropped after functionalization pass.
        view_3: f32[2, 32, 8, 8] = torch.ops.aten.view.default(getitem, [2, 32, 8, 8]);  getitem = None
        # ===========================================

        view_4: f32[1, 64, 8, 8] = torch.ops.aten.view.default(tangents_1, [1, 64, 8, 8]);  tangents_1 = None
        native_batch_norm_backward = torch.ops.aten.native_batch_norm_backward.default(view_4, view, None, repeat, repeat_1, getitem_1, getitem_2, True, 1e-05, [True, False, False]);  view_4 = view = repeat = repeat_1 = getitem_1 = getitem_2 = None
        getitem_3: f32[1, 64, 8, 8] = native_batch_norm_backward[0];  native_batch_norm_backward = None
        view_5: f32[2, 32, 8, 8] = torch.ops.aten.view.default(getitem_3, [2, 32, 8, 8]);  getitem_3 = None

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:21, code: out = out.relu()
        detach_3: f32[2, 32, 8, 8] = torch.ops.aten.detach.default(detach_2);  detach_2 = None
        detach_4: f32[2, 32, 8, 8] = torch.ops.aten.detach.default(detach_3);  detach_3 = None
        threshold_backward: f32[2, 32, 8, 8] = torch.ops.aten.threshold_backward.default(view_5, detach_4, 0);  view_5 = detach_4 = None

        # File: /home/rspring/workspace/debug_repro/dynamo_instance_norm.py:20, code: out = self.conv(inp)
        convolution_backward = torch.ops.aten.convolution_backward.default(threshold_backward, primals_5, primals_1, [0], [1, 1], [0, 0], [1, 1], False, [0, 0], 1, [False, True, False]);  threshold_backward = primals_5 = primals_1 = None
        getitem_7: f32[32, 32, 1, 1] = convolution_backward[1];  convolution_backward = None
        return pytree.tree_unflatten([view_3, getitem_7, None, None, None, None], self._out_spec)
```

Debug repro:
```python
import torch
import functorch
from torch._dynamo.optimizations.training import aot_nvprims_nvfuser, aot_nvprims_aten
from torch._prims.context import TorchRefsNvfuserCapabilityMode
optimize = torch._dynamo.optimize(aot_nvprims_aten)

torch.cuda.cudart().cudaProfilerStart()

class Fusion(torch.nn.Module):
  def __init__(self) :
    super(Fusion, self).__init__()
    self.conv = torch.nn.Conv2d(32, 32, (1, 1), bias=False)
    self.norm = torch.nn.InstanceNorm2d(32, track_running_stats=True)

  def forward(self, inp) :
    out = self.conv(inp)
    out = out.relu()
    out = self.norm(out)
    return out

model = Fusion().cuda()

input1 = torch.randn(2, 32, 8, 8, device="cuda")

optimized_model = optimize(model)
with torch.cuda.amp.autocast(False):
  for _ in range(5):
    out = optimized_model(input1)
    out.sum().backward()

torch.cuda.synchronize()
torch.cuda.cudart().cudaProfilerStop()
```


### Versions

upstream: a6acbad5

cc @ezyang @gchanan @zou3519 @mruberry @ngimel @Lezcano @fdrocha @peterbell10