[inductor][Seg fault] Inductor segfaulting with few AMP models
Labels: triaged, bug, oncall: pt2

### 🐛 Describe the bug

The following repro segfaults. Note that **it is flaky**, so might need to run multiple times.



### Error logs

~~~
Traceback (most recent call last):
  File "/scratch/anijain/work/pytorch/fail_triton.py", line 17, in <module>
    triton_fused_scatter_add_new_zeros_408_scatter_add_1_new_zeros_409_unsqueeze_1189_expand_15 = async_compile.triton('''
  File "/scratch/anijain/work/pytorch/torch/_inductor/codecache.py", line 463, in triton
    return _load_kernel(source_code)
  File "/scratch/anijain/work/pytorch/torch/_inductor/codecache.py", line 339, in _load_kernel
    kernel.precompile()
  File "/scratch/anijain/work/pytorch/torch/_inductor/triton_ops/autotune.py", line 58, in precompile
    self.launchers = [
  File "/scratch/anijain/work/pytorch/torch/_inductor/triton_ops/autotune.py", line 59, in <listcomp>
    self._precompile_config(c, warm_cache_only_with_cc)
  File "/scratch/anijain/work/pytorch/torch/_inductor/triton_ops/autotune.py", line 84, in _precompile_config
    binary = triton.compile(
  File "/scratch/anijain/work/triton/python/triton/compiler.py", line 1256, in compile
    asm, shared, kernel_name = _compile(fn, signature, device, constants, configs[0], num_warps, num_stages,
  File "/scratch/anijain/work/triton/python/triton/compiler.py", line 901, in _compile
    name, asm, shared_mem = _triton.code_gen.compile_ttir(backend, module, device, num_warps, num_stages, extern_libs, cc)
RuntimeError: Internal Triton PTX codegen error:
Segmentation fault (core dumped)
~~~

### Minified repro

~~~

from ctypes import c_void_p, c_long
import torch
import random
from torch import empty_strided, as_strided, device
from torch._inductor.codecache import AsyncCompile

aten = torch.ops.aten
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
async_compile = AsyncCompile()

import triton
import triton.language as tl
from torch._inductor.triton_ops.autotune import grid
from torch._C import _cuda_getCurrentRawStream as get_cuda_stream

triton_fused_scatter_add_new_zeros_408_scatter_add_1_new_zeros_409_unsqueeze_1189_expand_15 = async_compile.triton('''
import triton
import triton.language as tl
from torch._inductor.ir import ReductionHint
from torch._inductor.ir import TileHint
from torch._inductor.triton_ops.autotune import pointwise
from torch._inductor.utils import instance_descriptor

@pointwise(size_hints=[1048576], filename=__file__, meta={'signature': {0: '*i64', 1: '*fp16', 2: '*fp16', 3: 'i32'}, 'device': 0, 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2, 3), equal_to_1=())]})
@triton.jit
def kernel(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 900000
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK])
    xmask = xindex < xnumel
    x3 = (xindex // 90)
    x4 = xindex
    x0 = xindex % 90
    x2 = (xindex // 450000)
    tmp0 = tl.load(in_ptr0 + (x3), xmask)
    tmp1 = tl.load(in_ptr1 + (x4), xmask).to(tl.float32)
    tl.atomic_add(out_ptr0 + (x0 + (90*tmp0) + (6905250*x2) + tl.zeros([XBLOCK], tl.int32)), tmp1, xmask)
''')


async_compile.wait(globals())
del async_compile

~~~

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh