Repository: pytorch
Issue: Unsafe use of `at::parallel_for` in current codebase
Labels: high priority,module: performance,module: internals,module: autograd,triaged,module: multithreading

As discussed in https://github.com/pytorch/pytorch/issues/28370 , `parallel_for` is not propagating thread local states to worker threads which can lead to subtle errors when operations relying on thread local states are performed inside it.

Here is a list of all the potentially unsafe calls we have:
- [ ] In `Aten/native/Normalization.cpp` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/Normalization.cpp#L157-L180), [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/Normalization.cpp#L202-L236) and [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/Normalization.cpp#L275-L349) `.select()` is used.
- [ ] In `Aten/native/quantized/cpu/q_adaavgpool.cpp` [here](https://github.com/pytorch/pytorch/blob/eca01eb0a6f7f0f16c7e63a5b2c90ce91b7814a1/aten/src/ATen/native/quantized/cpu/q_adaavgpool.cpp#L168-L185) use of special ops on quantized tensor, not sure if they are safe?
- [ ] In `Aten/native/quantized/cpu/q_avgpool.cpp` [here](https://github.com/pytorch/pytorch/blob/eca01eb0a6f7f0f16c7e63a5b2c90ce91b7814a1/aten/src/ATen/native/quantized/cpu/q_avgpool.cpp#L36-L101), [here](https://github.com/pytorch/pytorch/blob/eca01eb0a6f7f0f16c7e63a5b2c90ce91b7814a1/aten/src/ATen/native/quantized/cpu/q_avgpool.cpp#L214-L235) and [here](https://github.com/pytorch/pytorch/blob/eca01eb0a6f7f0f16c7e63a5b2c90ce91b7814a1/aten/src/ATen/native/quantized/cpu/q_avgpool.cpp#L260-L280) use of `.contiguous(), special ops on quantized tensors and also nesting of parallel_for.
- [ ] In `ATen/native/quantized/cpu/qconv.cpp` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/quantized/cpu/qconv.cpp#L437-L495) Is fbgemm fine to use here? Also Tensor are unpacked inside the parallel_for block.
- [ ] In `ATen/native/quantized/cpu/qlinear.cpp` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/quantized/cpu/qlinear.cpp#L112) similar to conv above.
- [ ] In `ATen/native/SortingUtils.h` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/SortingUtils.h#L20-L40) `.select()` is used and the result passed to user function. User function called with Tensors so all user functions need to be checked.
- [ ] In `ATen/native/TensorIteratorReduce.cpp` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/TensorIteratorReduce.cpp#L44-L52), [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/TensorIteratorReduce.cpp#L151-L168), the second one is sharing of a TensorIterator, so not sure.
- [ ] In `ATen/native/Embedding.cpp` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/aten/src/ATen/native/Embedding.cpp#L105-L128) use of slicing and inplace op.
- [ ] In `binaries/intra_inter_benchmark.cc` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/binaries/intra_inter_benchmark.cc#L46-L64) we perform arbitrary ops on Tensors.
- [ ] In `csrc/api/include/torch/nn/parallel/data_parallel.h` [here](https://github.com/pytorch/pytorch/blob/42faf961c83d192933e199f8228909cc171765ca/torch/csrc/api/include/torch/nn/parallel/data_parallel.h#L207-L227) we forward a whole model.

cc @ezyang @gchanan @zou3519 @jerryzh168 @SsnL @albanD @gqchen @VitalyFedyunin @ngimel @mruberry @ifedan @ilia-cher 
What do you think is the way forward for this? Does any of you have some bandwidth to look into these and fix them?