wehub-resource-sync
175 lines
8.1 KiB
C++
175 lines
8.1 KiB
C++
// Copyright (c) 2022 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.
|
|
|
|
#include "paddle/phi/capi/include/c_kernel_registry.h"
|
|
|
|
#include "paddle/phi/capi/include/common.h"
|
|
#include "paddle/phi/capi/include/type_utils.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
void PD_KernelArgsParseFn(const phi::KernelKey& default_key,
|
|
phi::KernelArgsDef* args_def,
|
|
size_t in_nargs,
|
|
PD_KernelArgumentType* in_args_type,
|
|
size_t attr_nargs,
|
|
PD_KernelArgumentType* attr_args_type,
|
|
size_t out_nargs,
|
|
PD_KernelArgumentType* out_args_type) {
|
|
auto default_tensor_layout = phi::DataLayout::NCHW;
|
|
if (default_key.layout() != phi::DataLayout::ANY) {
|
|
default_tensor_layout = default_key.layout();
|
|
}
|
|
// inputs
|
|
for (size_t i = 0; i < in_nargs; ++i) {
|
|
auto arg_type = in_args_type[i];
|
|
if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_CONTEXT) {
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_TENSOR) {
|
|
args_def->AppendInput(default_key.backend(),
|
|
default_tensor_layout,
|
|
default_key.dtype(),
|
|
std::type_index(typeid(const phi::DenseTensor&)));
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_OPTIONAL_TENSOR) {
|
|
args_def->AppendInput(
|
|
default_key.backend(),
|
|
default_tensor_layout,
|
|
default_key.dtype(),
|
|
std::type_index(typeid(const paddle::optional<phi::DenseTensor>&)));
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_TENSOR) {
|
|
args_def->AppendInput(
|
|
default_key.backend(),
|
|
default_tensor_layout,
|
|
default_key.dtype(),
|
|
std::type_index(typeid(const std::vector<const phi::DenseTensor*>&)));
|
|
} else if (arg_type ==
|
|
PD_KernelArgumentType::PD_ARG_TYPE_OPTIONAL_MULTI_TENSOR) {
|
|
args_def->AppendInput(
|
|
default_key.backend(),
|
|
default_tensor_layout,
|
|
default_key.dtype(),
|
|
std::type_index(typeid(
|
|
const paddle::optional<std::vector<const phi::DenseTensor*>>&)));
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unavailable(
|
|
"PD_KernelArgumentType %d is not supported.", arg_type));
|
|
}
|
|
}
|
|
// attributes
|
|
for (size_t i = 0; i < attr_nargs; ++i) {
|
|
auto arg_type = attr_args_type[i];
|
|
if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_BOOL) {
|
|
args_def->AppendAttribute(phi::AttributeType::BOOL);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_FLOAT32) {
|
|
args_def->AppendAttribute(phi::AttributeType::FLOAT32);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_FLOAT64) {
|
|
args_def->AppendAttribute(phi::AttributeType::FLOAT64);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_INT32) {
|
|
args_def->AppendAttribute(phi::AttributeType::INT32);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_INT64) {
|
|
args_def->AppendAttribute(phi::AttributeType::INT64);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_STRING) {
|
|
args_def->AppendAttribute(phi::AttributeType::STRING);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_SCALAR) {
|
|
args_def->AppendAttribute(phi::AttributeType::SCALAR);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_INT_ARRAY) {
|
|
args_def->AppendAttribute(phi::AttributeType::INT_ARRAY);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_DATA_TYPE) {
|
|
args_def->AppendAttribute(phi::AttributeType::DATA_TYPE);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_DATA_LAYOUT) {
|
|
args_def->AppendAttribute(phi::AttributeType::DATA_LAYOUT);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_PLACE) {
|
|
args_def->AppendAttribute(phi::AttributeType::PLACE);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_BOOL) {
|
|
args_def->AppendAttribute(phi::AttributeType::BOOLS);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_INT32) {
|
|
args_def->AppendAttribute(phi::AttributeType::INT32S);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_INT64) {
|
|
args_def->AppendAttribute(phi::AttributeType::INT64S);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_FLOAT32) {
|
|
args_def->AppendAttribute(phi::AttributeType::FLOAT32S);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_FLOAT64) {
|
|
args_def->AppendAttribute(phi::AttributeType::FLOAT64S);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_STRING) {
|
|
args_def->AppendAttribute(phi::AttributeType::STRINGS);
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_SCALAR) {
|
|
args_def->AppendAttribute(phi::AttributeType::SCALARS);
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unavailable(
|
|
"PD_KernelArgumentType %d is not supported.", arg_type));
|
|
}
|
|
}
|
|
// outputs
|
|
for (size_t i = 0; i < out_nargs; ++i) {
|
|
auto arg_type = out_args_type[i];
|
|
if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_TENSOR) {
|
|
args_def->AppendOutput(default_key.backend(),
|
|
default_tensor_layout,
|
|
default_key.dtype(),
|
|
std::type_index(typeid(phi::DenseTensor*)));
|
|
} else if (arg_type == PD_KernelArgumentType::PD_ARG_TYPE_LIST_TENSOR) {
|
|
args_def->AppendOutput(
|
|
default_key.backend(),
|
|
default_tensor_layout,
|
|
default_key.dtype(),
|
|
std::type_index(typeid(std::vector<phi::DenseTensor*>)));
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unavailable(
|
|
"PD_KernelArgumentType %d is not supported.", arg_type));
|
|
}
|
|
}
|
|
}
|
|
|
|
void PD_RegisterPhiKernel(const char* kernel_name_cstr,
|
|
const char* backend_cstr,
|
|
PD_DataType pd_dtype,
|
|
PD_DataLayout pd_layout,
|
|
size_t in_nargs,
|
|
PD_KernelArgumentType* in_args_type,
|
|
size_t attr_nargs,
|
|
PD_KernelArgumentType* attr_args_type,
|
|
size_t out_nargs,
|
|
PD_KernelArgumentType* out_args_type,
|
|
void (*args_def_fn)(const PD_KernelKey*, PD_Kernel*),
|
|
void (*fn)(PD_KernelContext*),
|
|
void* variadic_kernel_fn) {
|
|
auto args_def_fn_wrapper = [args_def_fn](const phi::KernelKey& kernel_key,
|
|
phi::Kernel* kernel) {
|
|
args_def_fn(reinterpret_cast<const PD_KernelKey*>(&kernel_key),
|
|
reinterpret_cast<PD_Kernel*>(kernel));
|
|
};
|
|
phi::KernelFn kernel_fn = [fn](phi::KernelContext* ctx) {
|
|
fn(reinterpret_cast<PD_KernelContext*>(ctx));
|
|
};
|
|
std::string kernel_name(kernel_name_cstr);
|
|
|
|
auto dtype = phi::capi::ToPhiDataType(pd_dtype);
|
|
auto layout = phi::capi::ToPhiDataLayout(pd_layout);
|
|
phi::KernelKey kernel_key(
|
|
paddle::experimental::StringToBackend(backend_cstr), layout, dtype);
|
|
|
|
phi::Kernel kernel(kernel_fn, variadic_kernel_fn);
|
|
PD_KernelArgsParseFn(kernel_key,
|
|
kernel.mutable_args_def(),
|
|
in_nargs,
|
|
in_args_type,
|
|
attr_nargs,
|
|
attr_args_type,
|
|
out_nargs,
|
|
out_args_type);
|
|
|
|
args_def_fn_wrapper(kernel_key, &kernel);
|
|
phi::KernelFactory::Instance().kernels()[kernel_name][kernel_key] = kernel;
|
|
}
|
|
|
|
PD_REGISTER_CAPI(kernel_registry);
|