wehub-resource-sync
1064 lines
65 KiB
C++
1064 lines
65 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.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/phi/capi/include/common.h"
|
|
|
|
#if !defined(_WIN32)
|
|
|
|
namespace phi {
|
|
namespace capi {
|
|
|
|
#define CUSTOM_PHI_KERNEL(...) \
|
|
::phi::capi::CustomKernelImpl<decltype(&__VA_ARGS__), &__VA_ARGS__>::Compute
|
|
|
|
#define PHI_CAPI_INFER_META(...) \
|
|
::phi::capi::InferMetaFnImpl<decltype(&__VA_ARGS__), &__VA_ARGS__>::Call
|
|
|
|
#define CUSTOM_PHI_VARIADIC_KERNEL(...) \
|
|
reinterpret_cast<void *>( \
|
|
&::phi::capi::CustomKernelImpl<decltype(&__VA_ARGS__), \
|
|
&__VA_ARGS__>::VariadicCompute)
|
|
|
|
#define PD_CUSTOM_NARGS(...) \
|
|
_PD_CUSTOM_NARGS((__VA_ARGS__, _PD_CUSTOM_RESQ_N()))
|
|
#define _PD_CUSTOM_NARGS(...) _PD_CUSTOM_ARG_N(__VA_ARGS__)
|
|
#define _PD_CUSTOM_ARG_N_EXPAND( \
|
|
_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, N, ...) \
|
|
N
|
|
#define _PD_CUSTOM_ARG_N(args) _PD_CUSTOM_ARG_N_EXPAND args
|
|
#define _PD_CUSTOM_RESQ_N() 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0
|
|
|
|
#define PD_DATALAYOUT(arg__) PD_DataLayout::arg__
|
|
|
|
#ifdef __COUNTER__
|
|
#define PD_CUSTOM_PHI_KERNEL_ID __COUNTER__
|
|
#else
|
|
#define PD_CUSTOM_PHI_KERNEL_ID __LINE__
|
|
#endif
|
|
|
|
#define PD_CUSTOM_PHI_KERNEL_CONCATENATE(arg1, arg2) \
|
|
PD_CUSTOM_PHI_KERNEL_CONCATENATE1(arg1, arg2)
|
|
#define PD_CUSTOM_PHI_KERNEL_CONCATENATE1(arg1, arg2) \
|
|
PD_CUSTOM_PHI_KERNEL_CONCATENATE2(arg1, arg2)
|
|
#define PD_CUSTOM_PHI_KERNEL_CONCATENATE2(arg1, arg2) arg1##arg2
|
|
#define PD_CUSTOM_PHI_KERNEL_EXPAND(x) x
|
|
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION(N, meta_kernel_fn, backend, ...) \
|
|
PD_CUSTOM_PHI_KERNEL_CONCATENATE(_PD_BUILD_KERNEL_INSTANTIATION_, N) \
|
|
(meta_kernel_fn, backend, __VA_ARGS__)
|
|
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_1(meta_kernel_fn, backend, cpp_dtype) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_2( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_1(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_3( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_2(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_4( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_3(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_5( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_4(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_6( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_5(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_7( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_6(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_8( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_7(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_9( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_8(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_10( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_9(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_11( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_10(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_12( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_11(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_13( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_12(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_14( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_13(meta_kernel_fn, backend, __VA_ARGS__))
|
|
#define _PD_BUILD_KERNEL_INSTANTIATION_15( \
|
|
meta_kernel_fn, backend, cpp_dtype, ...) \
|
|
template decltype(meta_kernel_fn<cpp_dtype>) meta_kernel_fn<cpp_dtype>; \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_INSTANTIATION_14(meta_kernel_fn, backend, __VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_1(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
int TouchCustomKernelSymbolFor_##kernel_name##_##backend##_##layout() { \
|
|
return 0; \
|
|
}
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_2(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_1(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_3(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_2(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_4(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_3(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_5(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_4(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_6(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_5(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_7(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_6(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_8(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_7(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_9(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_8(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_10(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_9(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_11(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_10(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_12(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_11(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_13(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_12(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_14(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_13(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT_15(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
registrar_id, \
|
|
meta_kernel_fn, \
|
|
cpp_dtype, \
|
|
...) \
|
|
static const registrar_class<cpp_dtype> PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
__reg_phi_kernel_##kernel_name##_##backend##_##layout##_, registrar_id); \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_14(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define _PD_BUILD_KERNEL_REGISTRAR_INIT( \
|
|
N, registrar_class, kernel_name, backend, layout, meta_kernel_fn, ...) \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND(PD_CUSTOM_PHI_KERNEL_CONCATENATE( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT_, N)(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
PD_CUSTOM_PHI_KERNEL_ID, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define PD_BUILD_KERNEL_REGISTRAR_INIT( \
|
|
registrar_class, kernel_name, backend, layout, meta_kernel_fn, ...) \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND( \
|
|
_PD_BUILD_KERNEL_REGISTRAR_INIT(PD_CUSTOM_NARGS(__VA_ARGS__), \
|
|
registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define PD_BUILD_KERNEL_INSTANTIATION(meta_kernel_fn, backend, ...) \
|
|
_PD_BUILD_KERNEL_INSTANTIATION( \
|
|
PD_CUSTOM_NARGS(__VA_ARGS__), meta_kernel_fn, backend, __VA_ARGS__)
|
|
|
|
#define _PD_BUILD_2TA_KERNEL( \
|
|
registrar_class, kernel_name, backend, layout, meta_kernel_fn, ...) \
|
|
PD_BUILD_KERNEL_INSTANTIATION(meta_kernel_fn, backend, __VA_ARGS__); \
|
|
PD_BUILD_KERNEL_REGISTRAR_INIT(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__);
|
|
|
|
#define _PD_BUILD_PHI_KERNEL( \
|
|
registrar_class, kernel_name, backend, layout, meta_kernel_fn, ...) \
|
|
PD_CUSTOM_PHI_KERNEL_EXPAND(_PD_BUILD_2TA_KERNEL(registrar_class, \
|
|
kernel_name, \
|
|
backend, \
|
|
layout, \
|
|
meta_kernel_fn, \
|
|
__VA_ARGS__))
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_DEVICE_CONTEXT(dev_ctx) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<const dev_ctx &, Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(in_idx == 0, \
|
|
"Kernel's DeviceContext should appear before Inputs."); \
|
|
static_assert( \
|
|
attr_idx == 0, \
|
|
"Kernel's DeviceContext should appear before Attributes."); \
|
|
static_assert(out_idx == 0, \
|
|
"Kernel's DeviceContext should appear before Outputs."); \
|
|
dev_ctx arg = PD_GetDeviceContext(ctx); \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx + 1, in_idx, attr_idx, out_idx>( \
|
|
ctx, pargs..., arg); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
const dev_ctx &arg = std::get<idx>(ctx); \
|
|
auto dev_ctx_wrapper = phi::capi::DeviceContext( \
|
|
reinterpret_cast<PD_DeviceContext *>(const_cast<dev_ctx *>(&arg))); \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., dev_ctx_wrapper); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_INPUT(tensor_type) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<const tensor_type &, Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(attr_idx == 0, \
|
|
"Kernel's Input should appear before Attributes."); \
|
|
static_assert(out_idx == 0, \
|
|
"Kernel's Input should appear before Outputs."); \
|
|
const tensor_type arg = PD_InputAt(ctx, in_idx); \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx, in_idx + 1, attr_idx, out_idx>( \
|
|
ctx, pargs..., arg); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
const tensor_type &arg = std::get<idx>(ctx); \
|
|
auto tensor = phi::capi::DenseTensor( \
|
|
reinterpret_cast<PD_Tensor *>(const_cast<tensor_type *>(&arg))); \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., tensor); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_OPTIONAL_INPUT(tensor_type) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<const paddle::optional<tensor_type> &, \
|
|
Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(attr_idx == 0, \
|
|
"Kernel's Input should appear before Attributes."); \
|
|
static_assert(out_idx == 0, \
|
|
"Kernel's Input should appear before Outputs."); \
|
|
auto arg = PD_OptionalInputAt(ctx, in_idx); \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx, in_idx + 1, attr_idx, out_idx>( \
|
|
ctx, pargs..., arg); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
auto &arg = std::get<idx>(ctx); \
|
|
paddle::optional<tensor_type> tensor = \
|
|
PD_OptionalTensor(reinterpret_cast<PD_Tensor *>( \
|
|
const_cast<paddle::optional<tensor_type> *>(&arg))); \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., tensor); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_MULTI_INPUT(tensor_type) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<const std::vector<const tensor_type *> &, \
|
|
Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(attr_idx == 0, \
|
|
"Kernel's Input should appear before Attributes."); \
|
|
static_assert(out_idx == 0, \
|
|
"Kernel's Input should appear before Outputs."); \
|
|
auto arg = PD_MultiInputAt(ctx, in_idx); \
|
|
std::vector<const tensor_type *> tensor_ptr_vec; \
|
|
for (auto &tensor : arg) { \
|
|
tensor_ptr_vec.push_back(tensor.raw_data() ? &tensor : nullptr); \
|
|
} \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx, in_idx + 1, attr_idx, out_idx>( \
|
|
ctx, pargs..., tensor_ptr_vec); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
auto &arg = std::get<idx>(ctx); \
|
|
auto tensor_vec = PD_TensorVector(reinterpret_cast<PD_Tensor *>( \
|
|
const_cast<std::vector<const tensor_type *> *>(&arg))); \
|
|
std::vector<const tensor_type *> tensor_ptr_vec; \
|
|
for (auto &tensor : tensor_vec) { \
|
|
tensor_ptr_vec.push_back(tensor.raw_data() ? &tensor : nullptr); \
|
|
} \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., tensor_ptr_vec); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(attr_type) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<attr_type, Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(out_idx == 0, \
|
|
"Kernel's Attributes should appear before Outputs."); \
|
|
attr_type arg = PD_AttrAt<attr_type>(ctx, attr_idx); \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx, in_idx, attr_idx + 1, out_idx>( \
|
|
ctx, pargs..., arg); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
auto &arg = std::get<idx>(ctx); \
|
|
auto attr = PD_Attr<attr_type>(reinterpret_cast<void *>(&arg)); \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., attr); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF( \
|
|
attr_type) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<const attr_type &, Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(out_idx == 0, \
|
|
"Kernel's Attributes should appear before Outputs."); \
|
|
attr_type arg = PD_AttrAt<attr_type>(ctx, attr_idx); \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx, in_idx, attr_idx + 1, out_idx>( \
|
|
ctx, pargs..., arg); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
const attr_type &arg = std::get<idx>(ctx); \
|
|
auto attr = PD_Attr<attr_type>( \
|
|
reinterpret_cast<void *>(const_cast<attr_type *>(&arg))); \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., attr); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_OUTPUT(tensor_type) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<tensor_type *, Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
auto arg = PD_OutputAt(ctx, out_idx); \
|
|
tensor_type *ptr = (arg.raw_data() ? &arg : nullptr); \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx, in_idx, attr_idx, out_idx + 1>( \
|
|
ctx, pargs..., ptr); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
tensor_type *arg = std::get<idx>(ctx); \
|
|
auto tensor = \
|
|
phi::capi::DenseTensor(reinterpret_cast<PD_Tensor *>(arg)); \
|
|
auto tensor_ptr = tensor.raw_data() ? &tensor : nullptr; \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., tensor_ptr); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CustomKernelCallHelper_FOR_MULTI_OUTPUT(tensor_type) \
|
|
template <typename... Tail> \
|
|
struct CustomKernelCallHelper<std::vector<tensor_type *>, Tail...> { \
|
|
template <int dev_ctx_idx, \
|
|
int in_idx, \
|
|
int attr_idx, \
|
|
int out_idx, \
|
|
typename... PreviousArgs> \
|
|
static void Compute(PD_KernelContext *ctx, PreviousArgs &...pargs) { \
|
|
auto arg = PD_MultiOutputAt(ctx, out_idx); \
|
|
std::vector<tensor_type *> tensor_ptr_vec; \
|
|
for (auto &tensor : arg) { \
|
|
tensor_ptr_vec.push_back(tensor.raw_data() ? &tensor : nullptr); \
|
|
} \
|
|
CustomKernelCallHelper<Tail...>:: \
|
|
template Compute<dev_ctx_idx, in_idx, attr_idx, out_idx + 1>( \
|
|
ctx, pargs..., tensor_ptr_vec); \
|
|
} \
|
|
template <int idx, typename... PreviousArgs> \
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx, \
|
|
PreviousArgs &...pargs) { \
|
|
std::vector<tensor_type *> &arg = std::get<idx>(ctx); \
|
|
auto tensor_vec = PD_TensorVector(reinterpret_cast<PD_Tensor *>( \
|
|
const_cast<std::vector<tensor_type *> *>(&arg))); \
|
|
std::vector<tensor_type *> tensor_ptr_vec; \
|
|
for (auto &tensor : tensor_vec) { \
|
|
tensor_ptr_vec.push_back(tensor.raw_data() ? &tensor : nullptr); \
|
|
} \
|
|
return CustomKernelCallHelper<Tail...>::template VariadicCompute<idx + \
|
|
1>( \
|
|
ctx, pargs..., tensor_ptr_vec); \
|
|
} \
|
|
}
|
|
|
|
template <typename T>
|
|
struct CustomTypeTag {};
|
|
|
|
template <typename Fn, Fn fn>
|
|
struct CustomKernelImpl;
|
|
|
|
template <typename Return,
|
|
typename DevCtx,
|
|
typename... Args,
|
|
Return (*kernel_fn)(DevCtx, Args...)>
|
|
struct CustomKernelImpl<Return (*)(DevCtx, Args...), kernel_fn> {
|
|
static void Compute(PD_KernelContext *ctx) {
|
|
CustomKernelCallHelper<DevCtx, Args..., CustomTypeTag<int>>::
|
|
template Compute<0, 0, 0, 0>(ctx);
|
|
}
|
|
|
|
static void VariadicCompute(DevCtx dev_ctx, Args... args) {
|
|
const std::tuple<DevCtx, Args &...> args_tuple(dev_ctx, args...);
|
|
return CustomKernelCallHelper<DevCtx, Args..., CustomTypeTag<int>>::
|
|
template VariadicCompute<0>(args_tuple);
|
|
}
|
|
|
|
private:
|
|
template <typename... RemainingArgs>
|
|
struct CustomKernelCallHelper;
|
|
|
|
/* DeviceContext Helpers */
|
|
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_DEVICE_CONTEXT(
|
|
phi::capi::DeviceContext);
|
|
|
|
/* Input Helpers */
|
|
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_INPUT(phi::capi::DenseTensor);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_OPTIONAL_INPUT(
|
|
phi::capi::DenseTensor);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_MULTI_INPUT(phi::capi::DenseTensor);
|
|
|
|
/* Attribute Helpers */
|
|
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(bool);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(int32_t);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(int64_t);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(float);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(double);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(PD_DataType);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(PD_DataLayout);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_ATTRIBUTE(phi::capi::Place);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<bool>);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<int32_t>);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<int64_t>);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<float>);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<double>);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(std::string);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<std::string>);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
phi::capi::Scalar);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
phi::capi::IntArray);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<phi::capi::Scalar>);
|
|
|
|
/* Output Helpers */
|
|
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_OUTPUT(phi::capi::DenseTensor);
|
|
PD_SPECIALIZE_CustomKernelCallHelper_FOR_MULTI_OUTPUT(phi::capi::DenseTensor);
|
|
|
|
/* End case */
|
|
template <typename T>
|
|
struct CustomKernelCallHelper<CustomTypeTag<T>> {
|
|
template <int dev_ctx_idx, int in_idx, int attr_idx, int out_idx>
|
|
static void Compute(PD_KernelContext *ctx, DevCtx dev_ctx, Args &...args) {
|
|
static_assert(dev_ctx_idx > 0,
|
|
"Kernel should pass DeviceContext as argument.");
|
|
static_assert(out_idx > 0, "Kernel should have output argument.");
|
|
return kernel_fn(dev_ctx, args...);
|
|
}
|
|
|
|
template <int idx>
|
|
static void VariadicCompute(const std::tuple<DevCtx, Args &...> &ctx,
|
|
DevCtx dev_ctx,
|
|
Args... args) {
|
|
return kernel_fn(dev_ctx, args...);
|
|
}
|
|
};
|
|
};
|
|
|
|
template <typename Func>
|
|
struct CustomKernelArgsParseFunctor;
|
|
|
|
template <typename Return_, typename... Args_>
|
|
struct CustomKernelArgsParseFunctor<Return_ (*)(Args_...)> {
|
|
using Args = std::tuple<Args_...>;
|
|
enum : std::size_t { Arity = sizeof...(Args_) };
|
|
using Indices = std::make_index_sequence<Arity>;
|
|
template <std::size_t Index>
|
|
using Arg = typename std::tuple_element<Index, Args>::type;
|
|
|
|
CustomKernelArgsParseFunctor() {
|
|
auto args_type = ParseArgType(Indices{});
|
|
|
|
for (auto arg_type : args_type) {
|
|
if (arg_type ==
|
|
std::type_index(typeid(const phi::capi::DeviceContext *))) {
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const phi::capi::DenseTensor &))) {
|
|
in_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_TENSOR);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(
|
|
const paddle::optional<phi::capi::DenseTensor> &))) {
|
|
in_args_type.push_back(
|
|
PD_KernelArgumentType::PD_ARG_TYPE_OPTIONAL_TENSOR);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(
|
|
const std::vector<const phi::capi::DenseTensor *> &))) {
|
|
in_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_LIST_TENSOR);
|
|
} else if (arg_type ==
|
|
std::type_index(
|
|
typeid(const paddle::optional<
|
|
std::vector<const phi::capi::DenseTensor *>> &))) {
|
|
in_args_type.push_back(
|
|
PD_KernelArgumentType::PD_ARG_TYPE_OPTIONAL_MULTI_TENSOR);
|
|
} else if (arg_type == std::type_index(typeid(bool))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_BOOL);
|
|
} else if (arg_type == std::type_index(typeid(float))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_FLOAT32);
|
|
} else if (arg_type == std::type_index(typeid(double))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_FLOAT64);
|
|
} else if (arg_type == std::type_index(typeid(int32_t))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_INT32);
|
|
} else if (arg_type == std::type_index(typeid(int64_t))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_INT64);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const phi::capi::Place &))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_PLACE);
|
|
} else if (arg_type == std::type_index(typeid(const std::string &))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_STRING);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const std::vector<bool> &))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_LIST_BOOL);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const std::vector<float> &))) {
|
|
attr_args_type.push_back(
|
|
PD_KernelArgumentType::PD_ARG_TYPE_LIST_FLOAT32);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const std::vector<double> &))) {
|
|
attr_args_type.push_back(
|
|
PD_KernelArgumentType::PD_ARG_TYPE_LIST_FLOAT64);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const std::vector<int32_t> &))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_LIST_INT32);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const std::vector<int64_t> &))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_LIST_INT64);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const std::vector<std::string> &))) {
|
|
attr_args_type.push_back(
|
|
PD_KernelArgumentType::PD_ARG_TYPE_LIST_STRING);
|
|
} else if (arg_type == std::type_index(typeid(
|
|
const std::vector<phi::capi::Scalar> &))) {
|
|
attr_args_type.push_back(
|
|
PD_KernelArgumentType::PD_ARG_TYPE_LIST_SCALAR);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const phi::capi::Scalar &))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_SCALAR);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(const phi::capi::IntArray &))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_INT_ARRAY);
|
|
} else if (arg_type == std::type_index(typeid(PD_DataType))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_DATA_TYPE);
|
|
} else if (arg_type == std::type_index(typeid(PD_DataLayout))) {
|
|
attr_args_type.push_back(
|
|
PD_KernelArgumentType::PD_ARG_TYPE_DATA_LAYOUT);
|
|
} else if (arg_type == std::type_index(typeid(PD_DataLayout))) {
|
|
attr_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_PLACE);
|
|
} else if (arg_type ==
|
|
std::type_index(typeid(phi::capi::DenseTensor *))) {
|
|
out_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_TENSOR);
|
|
} else if (arg_type == std::type_index(typeid(
|
|
std::vector<phi::capi::DenseTensor *>))) {
|
|
out_args_type.push_back(PD_KernelArgumentType::PD_ARG_TYPE_LIST_TENSOR);
|
|
}
|
|
}
|
|
}
|
|
|
|
std::vector<PD_KernelArgumentType> in_args_type;
|
|
std::vector<PD_KernelArgumentType> attr_args_type;
|
|
std::vector<PD_KernelArgumentType> out_args_type;
|
|
|
|
private:
|
|
template <std::size_t... INDEX>
|
|
static std::vector<std::type_index> ParseArgType(
|
|
std::index_sequence<INDEX...>) {
|
|
return {std::type_index(typeid(Arg<INDEX>))...};
|
|
}
|
|
};
|
|
|
|
#define PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(attr_type) \
|
|
template <typename... Tail> \
|
|
struct InferMetaFnCallHelper<attr_type, Tail...> { \
|
|
template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs> \
|
|
static void Call(PD_InferMetaContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(out_idx == 0, \
|
|
"InferMeta's Attributes should appear before Outputs."); \
|
|
attr_type arg = PD_InferMetaAttrAt<attr_type>(ctx, attr_idx); \
|
|
InferMetaFnCallHelper< \
|
|
Tail...>::template Call<in_idx, attr_idx + 1, out_idx>(ctx, \
|
|
pargs..., \
|
|
arg); \
|
|
} \
|
|
}
|
|
|
|
#define PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF( \
|
|
attr_type) \
|
|
template <typename... Tail> \
|
|
struct InferMetaFnCallHelper<const attr_type &, Tail...> { \
|
|
template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs> \
|
|
static void Call(PD_InferMetaContext *ctx, PreviousArgs &...pargs) { \
|
|
static_assert(out_idx == 0, \
|
|
"InferMeta's Attributes should appear before Outputs."); \
|
|
attr_type arg = PD_InferMetaAttrAt<attr_type>(ctx, attr_idx); \
|
|
InferMetaFnCallHelper< \
|
|
Tail...>::template Call<in_idx, attr_idx + 1, out_idx>(ctx, \
|
|
pargs..., \
|
|
arg); \
|
|
} \
|
|
}
|
|
|
|
template <typename T>
|
|
struct InferMetaTypeTag {};
|
|
|
|
template <typename Fn, Fn fn>
|
|
struct InferMetaFnImpl;
|
|
|
|
template <typename Return, typename... Args, Return (*infer_meta_fn)(Args...)>
|
|
struct InferMetaFnImpl<Return (*)(Args...), infer_meta_fn> {
|
|
static void Call(PD_InferMetaContext *ctx) {
|
|
InferMetaFnCallHelper<Args...,
|
|
InferMetaTypeTag<int>>::template Call<0, 0, 0>(ctx);
|
|
}
|
|
|
|
private:
|
|
template <typename... RemainingArgs>
|
|
struct InferMetaFnCallHelper;
|
|
|
|
template <typename... Tail>
|
|
struct InferMetaFnCallHelper<const MetaTensor &, Tail...> {
|
|
template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
|
|
static void Call(PD_InferMetaContext *ctx, PreviousArgs &...pargs) {
|
|
static_assert(attr_idx == 0,
|
|
"InferMeta's Input should appear before Attributes.");
|
|
static_assert(out_idx == 0,
|
|
"InferMeta's Input should appear before Outputs.");
|
|
auto arg = MetaTensor(PD_InferMetaContextInputAt(ctx, in_idx));
|
|
InferMetaFnCallHelper<
|
|
Tail...>::template Call<in_idx + 1, attr_idx, out_idx>(ctx,
|
|
pargs...,
|
|
arg);
|
|
}
|
|
};
|
|
|
|
template <typename... Tail>
|
|
struct InferMetaFnCallHelper<const std::vector<const MetaTensor *> &,
|
|
Tail...> {
|
|
template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
|
|
static void Call(PD_InferMetaContext *ctx, PreviousArgs &...pargs) {
|
|
static_assert(attr_idx == 0,
|
|
"InferMeta's Input should appear before Attributes.");
|
|
static_assert(out_idx == 0,
|
|
"InferMeta's Input should appear before Outputs.");
|
|
auto arg = PD_InferMetaMultiInputAt(ctx, in_idx);
|
|
std::vector<const MetaTensor *> tensor_ptr_vec;
|
|
for (auto &tensor : arg) {
|
|
tensor_ptr_vec.push_back(tensor.raw_data() ? &tensor : nullptr);
|
|
}
|
|
InferMetaFnCallHelper<Tail...>::
|
|
template Call<in_idx + 1, attr_idx, out_idx>(
|
|
ctx, pargs..., tensor_ptr_vec);
|
|
}
|
|
};
|
|
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(bool);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(int);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(int64_t);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(float);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(DataType);
|
|
// PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(Backend);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_ATTRIBUTE(DataLayout);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(std::string);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(Scalar);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(IntArray);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<bool>);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<int>);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<int64_t>);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<float>);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<double>);
|
|
PD_SPECIALIZE_CAPI_InferMetaFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
|
|
std::vector<std::string>);
|
|
|
|
template <typename... Tail>
|
|
struct InferMetaFnCallHelper<MetaTensor *, Tail...> {
|
|
template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
|
|
static void Call(PD_InferMetaContext *ctx, PreviousArgs &...pargs) {
|
|
auto arg = MetaTensor(PD_InferMetaContextOutputAt(ctx, out_idx));
|
|
auto *arg_ptr = &arg;
|
|
InferMetaFnCallHelper<
|
|
Tail...>::template Call<in_idx, attr_idx, out_idx + 1>(ctx,
|
|
pargs...,
|
|
arg_ptr);
|
|
}
|
|
};
|
|
|
|
template <typename... Tail>
|
|
struct InferMetaFnCallHelper<std::vector<MetaTensor *>, Tail...> {
|
|
template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
|
|
static void Call(PD_InferMetaContext *ctx, PreviousArgs &...pargs) {
|
|
auto arg = PD_InferMetaMultiOutputAt(ctx, out_idx);
|
|
std::vector<MetaTensor *> tensor_ptr_vec;
|
|
for (auto &tensor : arg) {
|
|
tensor_ptr_vec.push_back(tensor.raw_data() ? &tensor : nullptr);
|
|
}
|
|
InferMetaFnCallHelper<Tail...>::
|
|
template Call<in_idx, attr_idx, out_idx + 1>(
|
|
ctx, pargs..., tensor_ptr_vec);
|
|
}
|
|
};
|
|
|
|
/* End case */
|
|
template <typename T>
|
|
struct InferMetaFnCallHelper<InferMetaTypeTag<T>> {
|
|
template <int in_idx, int attr_idx, int out_idx>
|
|
static void Call(PD_InferMetaContext *ctx, Args &...args) {
|
|
return infer_meta_fn(args...);
|
|
}
|
|
};
|
|
};
|
|
|
|
} // namespace capi
|
|
} // namespace phi
|
|
|
|
#endif
|