Repository: pytorch
Issue: crash@sleef_tryVXE2 () while trying to run torch.compile() BERT model
Labels: module: crash,triaged,module: sleef,module: POWER

### 🐛 Describe the bug

Hi, 

I have been trying to load a pre-trained bert model `bert-base-cased`  downloaded from HuggingFace and optimize it with torch.compile() on linux-s390x. It crashed - 

```python
0x000003ffb91c4418 in sleef_tryVXE2 () from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
(gdb) bt
#0  0x000003ffb91c4418 in sleef_tryVXE2 () from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#1  0x000003ffb91bedcc in cpuSupportsVXE2 () from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#2  0x000003ffb91c02c4 in disp_expf4_u10 () from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#3  0x000003ffb5dac610 in at::native::(anonymous namespace)::softmax_lastdim_kernel_impl(at::Tensor const&, at::Tensor const&) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#4  0x000003ffb4175890 in at::native::structured_softmax_cpu_out::impl(at::Tensor const&, long, bool, at::Tensor const&) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#5  0x000003ffb4e7ed92 in at::(anonymous namespace)::wrapper_CPU__softmax(at::Tensor const&, long, bool) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#6  0x000003ffb4e7ee28 in c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor (at::Tensor const&, long, bool), &at::(anonymous namespace)::wrapper_CPU__softmax>, at::Tensor, c10::guts::typelist::typelist<at::Tensor const&, long, bool> >, at::Tensor (at::Tensor const&, long, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, long, bool) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#7  0x000003ffb4b76af2 in at::_ops::_softmax::redispatch(c10::DispatchKeySet, at::Tensor const&, long, bool) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#8  0x000003ffb7949916 in torch::autograd::VariableType::(anonymous namespace)::_softmax(c10::DispatchKeySet, at::Tensor const&, long, bool) ()
  from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#9  0x000003ffb7949f8c in c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor (c10::DispatchKeySet, at::Tensor const&, long, bool), &torch::autograd::VariableType::(anonymous namespace)::_softmax>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor const&, long, bool> >, at::Tensor (c10::DispatchKeySet, at::Tensor const&, long, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, long, bool) () from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
...
..
arType>)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, long, std::optional<c10::ScalarType>) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#13 0x000003ffb4bf28b0 in at::_ops::softmax_int::call(at::Tensor const&, long, std::optional<c10::ScalarType>) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#14 0x000003ffbb39d020 in torch::autograd::THPVariable_softmax(_object*, _object*, _object*) ()
   from /root/anaconda3/envs/pytorch_test/lib/python3.10/site-packages/torch/lib/libtorch_python.so
#15 0x00000000010a7162 in method_vectorcall_VARARGS_KEYWORDS (func=<error reading variable: value has been optimized out>, args=args@entry=0x3ffb0ae43b0,
    nargsf=<optimized out>, kwnames=kwnames@entry=0x0) at /usr/local/src/conda/python-3.10.14/Objects/descrobject.c:344
```
**`Version:`**
Pytorch 2.2.0


**Reproduce:** 
Download pre-trained `bert-base-cased` from huggingface. https://huggingface.co/google-bert/bert-base-cased/tree/main

run the below scripts- 

```python 

from transformers import BertTokenizer, BertModel
import torch
tokenizer = BertTokenizer.from_pretrained('/root/bert-base-cased')
model = BertModel.from_pretrained("/root/bert-base-cased")
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)

# optimize with torch.compile

compile_model = torch.compile(model)
output = compile_model(**encoded_input)  -- it will crash here. 
```

### Versions

**Version:**
Pytorch 2.2.0
python  3.10.14
transformers  4.24.0
 
**platform** 
Linux-s390x