Repository: tensorflow
Issue: CopyTensor::ViaDMA function, allocator type sometimes not match actual input underlying memory type
Labels: stat:awaiting tensorflower,type:bug,comp:runtime,TF 2.12

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

source

### Tensorflow Version

2.12.0rc0

### Custom Code

Yes

### OS Platform and Distribution

CentOS Linux 7

### Mobile device

_No response_

### Python version

3.7.5

### Bazel version

bazel 3.7.2

### GCC/Compiler version

gcc-9

### CUDA/cuDNN version

cuda 11, cudnn 8

### GPU model and memory

Tesla V100S

### Current Behaviour?

In `CopyTensor::ViaDMA`, `alloc_attr `decides the direction of memory copy. However, sometimes `alloc_attr `does not keep the same as the Tensor pointer's underlying memory type. In my case,`src_alloc_attr.on_host()` is `True`, but `input->GetMemoryType()` equals to `kDevice`. So this results the memory copy direction in this function is cpu->gpu, but actually the direction should be gpu -> gpu. 
I think this bug does not reveal is because the cuda driver api, like `cuMemcpyHtoD()`, does not care about the direction if it's H to D or others, it only cares about the pointer attribute, if the src pointer is on device and dst pointer is also on device, even if we call `cuMemcpyHtoD()`, cuda driver would still do D to D copy. This feature would cover many bugs. 

I haven't figured out where did the `on_host `attribute is set. From my understanding so far, same allocator object would be reused on different tensors, but the `on_host `attribute is one-way, once it's been set `on_host`, it cannot be unset later. This might cause some issue? Also, why wouln't we just use `input->GetMemoryType()` to decieds the memory copy direction, instead of the `on_host `attribute of `alloc_attr`

I meet this issue when I run horovod unit test case. Add some log message in 
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/copy_tensor.cc#L219
, such as 
```
if(!src_alloc_attr.on_host() && (input->GetMemoryType()==AllocatorMemoryType::kHostPageable || input->GetMemoryType()==AllocatorMemoryType::kHostPinned)){
    std::cout<<"!!!!!!! src alloc not on host, but input mem type is on host"<< std::endl;
  }
  if( src_alloc_attr.on_host() && input->GetMemoryType()==AllocatorMemoryType::kDevice) {
    std::cout<<"!!!!!!! src alloc on host, but input mem is on device" << std::endl;
  }

  if(!dst_alloc_attr.on_host() && (output->GetMemoryType()==AllocatorMemoryType::kHostPageable || output->GetMemoryType()==AllocatorMemoryType::kHostPinned)){
    std::cout<<"!!!!!!! dst alloc not on host, but output mem type is on host"<< std::endl;
  }
  if( dst_alloc_attr.on_host() && output->GetMemoryType()==AllocatorMemoryType::kDevice) {
    std::cout<<"!!!!!!! dst alloc on host, but output mem is on device" << std::endl;
  }

```
For me, I ran horovod `alltoall Op` unit test case to reproduce this issue. But this issue might reveal in other cases.

### Standalone code to reproduce the issue

```shell
Run horovod unit test case can reproduce this issue:
https://github.com/horovod/horovod/blob/master/test/parallel/test_tensorflow.py
horovodrun --mpi -np 2 pytest -s -v test_tensorflow.py::TensorFlowTests::test_horovod_alltoall_gpu.
```


### Relevant log output

_No response_</details>