wehub-resource-sync
263 lines
8.2 KiB
C++
263 lines
8.2 KiB
C++
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License. */
|
|
|
|
#pragma once
|
|
|
|
#include <glog/logging.h>
|
|
#include <limits>
|
|
#include <string>
|
|
#include <utility>
|
|
#include "paddle/common/layout.h"
|
|
#include "paddle/phi/api/include/tensor.h"
|
|
#include "paddle/phi/api/lib/backend_set.h"
|
|
#include "paddle/phi/api/lib/data_type_set.h"
|
|
#include "paddle/phi/backends/all_context.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
|
|
#include "paddle/phi/core/selected_rows.h"
|
|
#include "paddle/phi/core/sparse_coo_tensor.h"
|
|
#include "paddle/phi/core/sparse_csr_tensor.h"
|
|
|
|
// TODO(chenweihang): split Key, Kernel, Factory into diff files
|
|
#include "paddle/phi/core/kernel_factory.h"
|
|
|
|
namespace paddle {
|
|
namespace experimental {
|
|
|
|
namespace detail {
|
|
PADDLE_API BackendSet GetTensorBackendSet(const phi::TensorBase& t);
|
|
PADDLE_API std::size_t CountLeadingZeros(uint32_t val);
|
|
} // namespace detail
|
|
|
|
PADDLE_API phi::DeviceContext* GetDeviceContextByBackend(phi::Backend backend);
|
|
|
|
enum class KernelType {
|
|
DENSE_TENSOR_KERNEL, // kernel for DenseTensor
|
|
SELECTED_ROWS_KERNEL, // kernel for SelectedRows
|
|
SPARSE_COO_KERNEL, // kernel for SparseCooTensor
|
|
SPARSE_CSR_KERNEL // kernel for SparseCsrTensor
|
|
};
|
|
|
|
// TODO(chenweihang): support DataLayout and DataType selected
|
|
struct KernelKeySet {
|
|
BackendSet backend_set{Backend::UNDEFINED};
|
|
phi::DataLayout layout{phi::DataLayout::UNDEFINED};
|
|
DataType dtype{DataType::UNDEFINED};
|
|
|
|
// TODO(chenweihang): iterate all kernel key for kernel selection
|
|
phi::KernelKey GetHighestPriorityKernelKey() {
|
|
uint32_t bitset_value = backend_set.bitset();
|
|
#ifdef PADDLE_WITH_CUSTOM_DEVICE
|
|
if (backend_set.Has(Backend(4))) {
|
|
return phi::KernelKey(Backend(4), layout, dtype);
|
|
}
|
|
#endif
|
|
std::size_t leading_zeros = detail::CountLeadingZeros(bitset_value);
|
|
Backend selected_backend = static_cast<Backend>(32 - leading_zeros);
|
|
VLOG(8) << "GetHighestPriorityKernelKey: selected_backend = "
|
|
<< selected_backend;
|
|
|
|
return phi::KernelKey(selected_backend, layout, dtype);
|
|
}
|
|
};
|
|
|
|
namespace detail {
|
|
|
|
template <typename Functor>
|
|
struct ArgsIterator {
|
|
template <typename... Args>
|
|
inline Functor& apply() {
|
|
return self();
|
|
}
|
|
|
|
template <typename T, typename... Args>
|
|
inline Functor& apply(T&& arg, Args&&... args) {
|
|
self()(std::forward<T>(arg));
|
|
if (self().short_circuit()) {
|
|
return self();
|
|
} else {
|
|
return apply(std::forward<Args>(args)...);
|
|
}
|
|
}
|
|
|
|
constexpr bool short_circuit() const { return false; }
|
|
|
|
private:
|
|
inline Functor& self() { return *static_cast<Functor*>(this); }
|
|
};
|
|
|
|
struct KernelKeyParser : ArgsIterator<KernelKeyParser> {
|
|
KernelKeySet key_set;
|
|
bool disable_gpudnn = false;
|
|
// this dtype_set is used for cache multi-inputs dtype and used for
|
|
// data_promote
|
|
DataTypeSet dtype_set{DataType::UNDEFINED};
|
|
|
|
inline void AssignKernelKeySet(const phi::TensorBase& tensor) {
|
|
// assign Backend
|
|
BackendSet tensor_backend_set = detail::GetTensorBackendSet(tensor);
|
|
key_set.backend_set = key_set.backend_set | tensor_backend_set;
|
|
// tensor's attribute use_gpudnn=False, explicitly disable gpudnn kernel
|
|
if (tensor_backend_set ==
|
|
BackendSet(paddle::experimental::get_accelerat_backend()) ||
|
|
disable_gpudnn) {
|
|
disable_gpudnn = true;
|
|
key_set.backend_set = key_set.backend_set - BackendSet(Backend::GPUDNN);
|
|
VLOG(8) << "Disable kernel backend: GPUDNN";
|
|
}
|
|
// assign DataLayout
|
|
phi::DataLayout tensor_layout = tensor.layout();
|
|
key_set.layout =
|
|
tensor_layout > key_set.layout ? tensor_layout : key_set.layout;
|
|
// assign DataType
|
|
key_set.dtype = tensor.dtype();
|
|
dtype_set = dtype_set | DataTypeSet(key_set.dtype);
|
|
auto promote_result = PromoteTypes(dtype_set);
|
|
if (promote_result != DataType::UNDEFINED) {
|
|
key_set.dtype = promote_result;
|
|
VLOG(8) << "promote kernel DataType:" << promote_result;
|
|
}
|
|
}
|
|
|
|
void operator()(const Tensor& x) {
|
|
const auto* tensor = x.impl().get();
|
|
if (tensor) {
|
|
AssignKernelKeySet(*tensor);
|
|
}
|
|
}
|
|
|
|
void operator()(const std::vector<Tensor>& x) {
|
|
if (!x.empty()) {
|
|
const phi::TensorBase& tensor = *x.at(0).impl();
|
|
AssignKernelKeySet(tensor);
|
|
}
|
|
}
|
|
|
|
void operator()(const paddle::optional<Tensor>& x) {
|
|
if (x) {
|
|
const phi::TensorBase& tensor = *(x.get_ptr()->impl());
|
|
AssignKernelKeySet(tensor);
|
|
}
|
|
}
|
|
|
|
// skip other type args, these args don't used in kernel selection
|
|
template <typename T>
|
|
void operator()(const T& x) {
|
|
// do nothing
|
|
}
|
|
};
|
|
|
|
struct KernelTypeParser : ArgsIterator<KernelTypeParser> {
|
|
KernelType kernel_type{KernelType::DENSE_TENSOR_KERNEL};
|
|
|
|
// TODO(chenweihang): deal with multiple diff input Tensors
|
|
// TODO(chenweihang): add global device guard method to set backend
|
|
void operator()(const Tensor& x) {
|
|
if (phi::SelectedRows::classof(x.impl().get())) {
|
|
kernel_type = KernelType::SELECTED_ROWS_KERNEL;
|
|
} else if (phi::SparseCooTensor::classof(x.impl().get())) {
|
|
kernel_type = KernelType::SPARSE_COO_KERNEL;
|
|
} else if (phi::SparseCsrTensor::classof(x.impl().get())) {
|
|
kernel_type = KernelType::SPARSE_CSR_KERNEL;
|
|
}
|
|
}
|
|
|
|
// skip other type args, these args don't used in kernel selection
|
|
template <typename T>
|
|
void operator()(const T& x) {
|
|
// do nothing
|
|
}
|
|
};
|
|
|
|
/* ------------------ for auto parallel ----------------------- */
|
|
|
|
struct DistTensorTypeParser : ArgsIterator<DistTensorTypeParser> {
|
|
bool result = false;
|
|
|
|
bool short_circuit() { return result; }
|
|
|
|
void operator()(const Tensor& x) { result = x.is_dist_tensor(); }
|
|
|
|
void operator()(const paddle::optional<Tensor>& x) {
|
|
if (x) {
|
|
result = x.get_ptr()->is_dist_tensor();
|
|
}
|
|
}
|
|
|
|
void operator()(const std::vector<Tensor>& x) {
|
|
if (!x.empty()) {
|
|
for (auto& t : x) {
|
|
result = result || t.is_dist_tensor();
|
|
if (short_circuit()) break;
|
|
}
|
|
}
|
|
}
|
|
|
|
void operator()(const paddle::optional<std::vector<Tensor>>& x) {
|
|
if (x && !x->empty()) {
|
|
for (auto& t : *(x.get_ptr())) {
|
|
result = result || t.is_dist_tensor();
|
|
if (short_circuit()) break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// skip other type args, these args don't used in kernel selection
|
|
template <typename T>
|
|
void operator()(const T& x) {
|
|
// do nothing
|
|
}
|
|
};
|
|
|
|
} // namespace detail
|
|
|
|
template <typename... Args>
|
|
KernelKeySet ParseKernelKeyByInputArgs(const Args&... args) {
|
|
return detail::KernelKeyParser().apply(args...).key_set;
|
|
}
|
|
|
|
template <typename... Args>
|
|
KernelType ParseKernelTypeByInputArgs(const Args&... args) {
|
|
return detail::KernelTypeParser().apply(args...).kernel_type;
|
|
}
|
|
|
|
DataType ParseDataType(DataType dtype);
|
|
DataType ParseDataType(const Tensor& tensor);
|
|
DataType ParseDataType(const std::vector<Tensor>& tensors);
|
|
DataType ParseDataTypeWithInputOrder(DataType dtype, const Tensor& tensor);
|
|
|
|
PADDLE_API Backend ParseBackend(const Place& place);
|
|
Backend ParseBackend(const Tensor& tensor);
|
|
template <typename T, typename... Args>
|
|
Backend ParseBackend(T t, Args... args) {
|
|
auto backend_set =
|
|
BackendSet(ParseBackend(t)) | BackendSet(ParseBackend(args...));
|
|
return static_cast<Backend>(32 -
|
|
detail::CountLeadingZeros(backend_set.bitset()));
|
|
}
|
|
Backend ParseBackendWithInputOrder(const Place& place, const Tensor& tensor);
|
|
|
|
PADDLE_API phi::DataLayout ParseLayout(phi::DataLayout layout);
|
|
phi::DataLayout ParseLayout(const Tensor& tensor);
|
|
phi::DataLayout ParseLayoutWithInputOrder(phi::DataLayout layout,
|
|
const Tensor& tensor);
|
|
|
|
template <typename... Args>
|
|
bool AllInputsAreDistTensor(const Args&... args) {
|
|
return detail::DistTensorTypeParser().apply(args...).result;
|
|
}
|
|
|
|
} // namespace experimental
|
|
} // namespace paddle
|