| 171 | #endif |
| 172 | |
| 173 | inline phi::DenseTensor TransDataType(const phi::DenseTensor& tensor, |
| 174 | DataType dtype) { |
| 175 | auto& pool = phi::DeviceContextPool::Instance(); |
| 176 | |
| 177 | VLOG(3) << "DataTypeTransform src_dtype: " << tensor.dtype() |
| 178 | << " dst_dtype: " << dtype; |
| 179 | |
| 180 | DefaultAllocator alloc(tensor.place()); |
| 181 | phi::DenseTensor out(&alloc, {dtype, tensor.dims(), tensor.layout()}); |
| 182 | |
| 183 | if (tensor.place().GetType() == phi::AllocationType::CPU) { |
| 184 | auto* dev_ctx = static_cast<phi::CPUContext*>(pool.Get(tensor.place())); |
| 185 | return CastDataType(*dev_ctx, tensor, dtype); |
| 186 | #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) |
| 187 | } else if (tensor.place().GetType() == phi::AllocationType::GPU) { |
| 188 | auto* dev_ctx = static_cast<phi::GPUContext*>(pool.Get(tensor.place())); |
| 189 | return CastDataType(*dev_ctx, tensor, dtype); |
| 190 | #endif |
| 191 | #ifdef PADDLE_WITH_XPU |
| 192 | } else if (tensor.place().GetType() == phi::AllocationType::XPU) { |
| 193 | auto* dev_ctx = static_cast<phi::XPUContext*>(pool.Get(tensor.place())); |
| 194 | return CastDataType(*dev_ctx, tensor, dtype); |
| 195 | #endif |
| 196 | #ifdef PADDLE_WITH_CUSTOM_DEVICE |
| 197 | } else if (tensor.place().GetType() == phi::AllocationType::CUSTOM) { |
| 198 | phi::DenseTensor out; |
| 199 | out.Resize(tensor.dims()); |
| 200 | auto* dev_ctx = static_cast<phi::CustomContext*>(pool.Get(tensor.place())); |
| 201 | auto kernel_result = |
| 202 | phi::KernelFactory::Instance().SelectKernelOrThrowError( |
| 203 | "cast", |
| 204 | {phi::TransToPhiBackend(tensor.place()), |
| 205 | phi::DataLayout::ALL_LAYOUT, |
| 206 | tensor.dtype()}); |
| 207 | using kernel_signature = void (*)(const phi::DeviceContext&, |
| 208 | const phi::DenseTensor&, |
| 209 | phi::DataType, |
| 210 | phi::DenseTensor*); |
| 211 | const auto& kernel = kernel_result.kernel; |
| 212 | auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>(); |
| 213 | (*kernel_fn)(*dev_ctx, tensor, dtype, &out); |
| 214 | return out; |
| 215 | #endif |
| 216 | } else { |
| 217 | PADDLE_THROW(common::errors::Unimplemented( |
| 218 | "Place type is not supported when casting data type.")); |
| 219 | } |
| 220 | return out; |
| 221 | } |
| 222 | |
| 223 | inline phi::DenseTensor TransDataPlace(const phi::DenseTensor& tensor, |
| 224 | Place dst_place) { |