| 46 | ////////////////// Forward api impls ////////////////////// |
| 47 | |
| 48 | Tensor add_n_impl(const std::vector<Tensor>& x) { |
| 49 | Backend kernel_backend = Backend::UNDEFINED; |
| 50 | DataLayout kernel_layout = DataLayout::UNDEFINED; |
| 51 | DataType kernel_data_type = DataType::UNDEFINED; |
| 52 | |
| 53 | if (kernel_backend == Backend::UNDEFINED || |
| 54 | kernel_layout == DataLayout::UNDEFINED || |
| 55 | kernel_data_type == DataType::UNDEFINED) { |
| 56 | auto kernel_key_set = ParseKernelKeyByInputArgs(x); |
| 57 | auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey(); |
| 58 | if (kernel_backend == Backend::UNDEFINED) { |
| 59 | kernel_backend = kernel_key.backend(); |
| 60 | } |
| 61 | if (kernel_layout == DataLayout::UNDEFINED) { |
| 62 | kernel_layout = kernel_key.layout(); |
| 63 | } |
| 64 | if (kernel_data_type == DataType::UNDEFINED) { |
| 65 | kernel_data_type = kernel_key.dtype(); |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | bool is_sr_kernel = true; |
| 70 | for (auto& input : x) { |
| 71 | if (phi::DenseTensor::classof(input.impl().get()) || |
| 72 | phi::distributed::DistTensor::classof(input.impl().get())) { |
| 73 | is_sr_kernel = false; |
| 74 | break; |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | const std::string kernel_name = (is_sr_kernel ? "add_n_sr" : "add_n"); |
| 79 | |
| 80 | VLOG(6) << "add_n API kernel key: [" << kernel_backend << ", " |
| 81 | << kernel_layout << ", " << kernel_data_type << "]"; |
| 82 | auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError( |
| 83 | kernel_name, {kernel_backend, kernel_layout, kernel_data_type}); |
| 84 | const auto& kernel = kernel_result.kernel; |
| 85 | VLOG(6) << kernel_name << " kernel: " << kernel; |
| 86 | auto* dev_ctx = GetDeviceContextByBackend( |
| 87 | kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend); |
| 88 | |
| 89 | Tensor api_output; |
| 90 | |
| 91 | if (is_sr_kernel) { |
| 92 | std::vector<const phi::SelectedRows*> input_x(x.size()); |
| 93 | for (size_t i = 0; i < input_x.size(); ++i) { |
| 94 | input_x[i] = static_cast<phi::SelectedRows*>(x[i].impl().get()); |
| 95 | } |
| 96 | auto x_meta_vec = MakeMetaTensor(input_x); |
| 97 | std::vector<const phi::MetaTensor*> x_metas(x_meta_vec.size()); |
| 98 | for (size_t i = 0; i < x_meta_vec.size(); ++i) { |
| 99 | x_metas[i] = &x_meta_vec[i]; |
| 100 | } |
| 101 | auto kernel_out = SetSelectedRowsKernelOutput(&api_output); |
| 102 | phi::MetaTensor meta_out(kernel_out); |
| 103 | phi::AddNInferMeta(x_metas, &meta_out); |
| 104 | |
| 105 | using kernel_signature = |
nothing calls this directly
no test coverage detected