Register functorch XPU/HPU dispatch keys (#167095)

Fixes TestOperatorsXPU.test_data_write_errors_under_transform_xpu https://github.com/intel/torch-xpu-ops/issues/2237

Tests on other devices throw runtime error "_mutating directly with `.data` inside functorch transform is not allowed._", but XPU/HPU fails earlier on `_has_compatible_shallow_copy_type`. This check is not met only when calling tensor.data inside functorch call.

```cpp
bool _has_compatible_shallow_copy_type(const Tensor& self, const Tensor& from) {
  return self.unsafeGetTensorImpl()->has_compatible_shallow_copy_type(
      from.key_set());
}
```

### t.data
| Tensor | Device | Dispatch Keys |
|--------|---------|---------------|
| `self` | `xpu` | `XPU, ADInplaceOrView, AutogradXPU, AutocastXPU` |
| `from` | `cpu` | `CPU, ADInplaceOrView, AutogradCPU, AutocastCPU` |

### t.data inside functorch transform
| Tensor | Device | Dispatch Keys |
|--------|---------|---------------|
| `self` | `xpu` | `ADInplaceOrView, AutogradOther, FuncTorchGradWrapper` |
| `from` | `cpu` | `CPU, ADInplaceOrView, AutogradCPU, AutocastCPU, FuncTorchGradWrapper` |

### t.data inside functorch transform + XPU dispatch key
| Tensor | Device | Dispatch Keys |
|--------|---------|---------------|
| `self` | `xpu` | `XPU, ADInplaceOrView, AutogradXPU, AutocastXPU, FuncTorchGradWrapper` |
| `from` | `cpu` | `CPU, ADInplaceOrView, AutogradCPU, AutocastCPU, FuncTorchGradWrapper` |
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167095
Approved by: https://github.com/guangyey, https://github.com/albanD
This commit is contained in:
Slawomir Siwek
2025-11-10 03:10:22 +00:00
committed by PyTorch MergeBot
parent fe6615e397
commit 2c78080ec0

View File

@ -157,6 +157,8 @@ constexpr DispatchKeySet kKeysToPropagateToWrapper({
DispatchKey::Negative,
DispatchKey::Conjugate,
DispatchKey::XLA,
DispatchKey::XPU,
DispatchKey::HPU,
DispatchKey::CUDA,
DispatchKey::CPU,
DispatchKey::PrivateUse1,