| 987 | } |
| 988 | |
| 989 | std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor( |
| 990 | std::shared_ptr<phi::distributed::DistTensor> input, |
| 991 | const phi::TensorArgDef& target_args_def, |
| 992 | const TransformFlag& transform_flag, |
| 993 | bool is_stride_kernel) { |
| 994 | if (input) { |
| 995 | phi::distributed::DistTensor* dist_tensor = input.get(); |
| 996 | const phi::DenseTensor& dense_tensor = dist_tensor->value(); |
| 997 | if (!transform_flag.NeedTransform() || !dense_tensor.initialized() || |
| 998 | (!NeedTransformPlace( |
| 999 | dense_tensor.place(), target_args_def.backend, transform_flag) && |
| 1000 | !NeedTransformDataType( |
| 1001 | dense_tensor.dtype(), target_args_def.dtype, transform_flag) && |
| 1002 | !NeedTransformLayout(dense_tensor.layout(), |
| 1003 | target_args_def.layout, |
| 1004 | dense_tensor.place(), |
| 1005 | transform_flag) && |
| 1006 | !NeedTransform2Contiguous(is_stride_kernel, |
| 1007 | dense_tensor.meta().is_contiguous()))) { |
| 1008 | if (NeedTransform2Contiguous(is_stride_kernel, |
| 1009 | dense_tensor.meta().is_contiguous()) && |
| 1010 | dense_tensor.initialized()) { |
| 1011 | auto dist_out = std::make_shared<phi::distributed::DistTensor>( |
| 1012 | dist_tensor->dims(), dist_tensor->dist_attr()); |
| 1013 | auto* out = dist_out->unsafe_mutable_value(); |
| 1014 | *out = Trans2Contiguous(dense_tensor); |
| 1015 | return dist_out; |
| 1016 | } |
| 1017 | return input; |
| 1018 | } |
| 1019 | // TODO(chenweihang): The global meta in DistTensor is not changed, |
| 1020 | // but the local meta in DenseTensor maybe changed, such as layout |
| 1021 | // change(NCHW->NHWC), so the new DistTensor's meta maybe not unified. |
| 1022 | VLOG(6) << "PrepareDataForDistTensor return transformed dist tensor"; |
| 1023 | auto dist_out = std::make_shared<phi::distributed::DistTensor>( |
| 1024 | dist_tensor->dims(), dist_tensor->dist_attr()); |
| 1025 | auto* out = dist_out->unsafe_mutable_value(); |
| 1026 | *out = TransformData( |
| 1027 | dense_tensor, target_args_def, transform_flag, is_stride_kernel); |
| 1028 | return dist_out; |
| 1029 | } |
| 1030 | return nullptr; |
| 1031 | } |
| 1032 | |
| 1033 | std::vector<std::shared_ptr<phi::distributed::DistTensor>> |
| 1034 | PrepareDataForDistTensor( |
no test coverage detected