Save mutated arguments after the kernel call in generated autograd kernels - #193758

…rnels

aten::rrelu_with_noise takes `noise` as a mutable output argument: the kernel
writes the sampled negative slopes into it, and rrelu_with_noise_backward scales
the incoming gradient with them. The generated autograd kernel built
SavedVariable(noise, false) in the setup_derivative block, before the kernel
call, so the packed value was the buffer's pre-call content.

Without saved tensor hooks that is harmless, because the SavedVariable holds on
to the argument and therefore aliases the buffer the kernel fills. With a hook
that materializes at pack time it is not. save_on_cpu copies uninitialized
memory. Non-reentrant checkpointing is worse: during recomputation the pack for
`noise` raises _StopRecomputationError once the saved-tensor counter reaches the
count from the original forward, so when `noise` is among the last saved tensors
of a segment the replay is aborted before the kernel runs and the buffer is
never filled. Backward then scales gradients by uninitialized memory while the
forward output stays bit-exact, so no RNG or output probe can see it. Because
rrelu_with_noise is SchemaKind.mutable rather than inplace/out, no
increment_version is emitted for `noise` and SavedVariable's staleness check
cannot catch it either.

Save arguments that the operator mutates after the call instead. That is the
value backward already reads today through aliasing, so the no-hook path is
unchanged. The filter uses post_self_positional_mutable, which excludes the
`self` argument of in-place ops; that one is already handled by original_self
and the saved outputs.

The derivative for the autogenerated rrelu_with_noise_functional also referred
to the `noise` input, which that op never writes: it clones the argument and
returns the filled clone as `noise_out`. Changed to use noise_out.

Besides the two rrelu ops, this moves the save point for running_mean and
running_var in _native_batch_norm_legit and _batch_norm_with_update, the only
other ops in derivatives.yaml with a saved mutated argument. Their backward
reads the running stats only when training=False, and the kernel does not update
them in that case, so there is no numerical change; the no-hook path already
observed the post-update values.

I considered the narrower fix of routing at::native::rrelu through
rrelu_with_noise_functional instead. It does not cover nn.RReLU(inplace=True),
which goes through rrelu_ -> rrelu_with_noise_ and has the same problem, without
adding an extra copy, so I fixed the save point in codegen instead.

The new tests pass a sentinel value as `noise` rather than relying on F.rrelu's
empty_like buffer. A stale read of that buffer is uninitialized memory, which can
happen to hold the correct slopes when the allocator hands back the block the
reference run just freed, so the sentinel is what makes the regression
deterministic. With it, all four failing tests report the same
6.898226737976074 maximum deviation on every run.

Test Plan:

```
python test/test_nn.py -k rrelu
python test/test_nn.py -k batchnorm
python test/test_autograd.py -k saved_tensor
python test/test_autograd.py -k checkpoint
python test/functorch/test_aotdispatch.py -k rrelu
python test/test_ops.py -k rrelu
python test/test_ops_gradients.py -k rrelu
python test/test_modules.py -k BatchNorm
python test/test_decomp.py -k batch_norm
```

All pass on a CPU build, with one exception that is not related to this change:
test_batchnorm_nhwc_cpu fails identically with and without it, down to the same
2.3365020751953125e-05 deviation at index 0. It compares channels_last against
contiguous BatchNorm3d gradients, and precisons[torch.float16] is None, so the
fp16 case runs at the float32 default tolerance while float32 and bfloat16 get a
relaxed 1e-4.

test_rrelu_saved_noise_hooks (both inplace values),
test_rrelu_saved_noise_non_reentrant_checkpoint with early_stop=True, and
test_rrelu_with_noise_functional_backward fail before this change and pass after.
The early_stop=False case passes either way and is there as a control: it is the
configuration where the replay reaches the kernel.

Fixes pytorch#193671

This change was prepared with the assistance of an AI coding assistant; the
analysis, code and tests were reviewed by the author.