Paddle Custom device的仓库位于:github.com/PaddlePaddl... 其示例代码为在backends/custom_cpu,本文主要讲解一下custom_cpu使用的paddle相关的代码,介绍paddle是如何实现custom device的。
1 Custom cpu算子注册代码
以add算子为例,其注册代码为于backends/custom_cpu/kernels/elementwise_kernel.cc
,代码如下:
cpp
template <typename T>
void AddKernel(const phi::Context& dev_ctx,
const phi::DenseTensor& x,
const phi::DenseTensor& y,
phi::DenseTensor* out) {
int axis = -1;
custom_kernel::AddRawKernel<T>(dev_ctx, x, y, axis, out);
}
PD_BUILD_PHI_KERNEL(add,
custom_cpu,
ALL_LAYOUT,
custom_kernel::AddKernel,
int32_t,
int64_t,
float,
double) {}
注意,该文件include的头文件中,和paddle相关的只有paddle/phi/capi/all.h
,该文件内容如下:
h
#pragma once
#if !defined(_WIN32)
#include "paddle/phi/capi/include/c_data_type.h"
#include "paddle/phi/capi/include/c_device_context.h"
#include "paddle/phi/capi/include/c_infer_meta_context.h"
#include "paddle/phi/capi/include/c_int_array.h"
#include "paddle/phi/capi/include/c_kernel_context.h"
#include "paddle/phi/capi/include/c_kernel_factory.h"
#include "paddle/phi/capi/include/c_kernel_registry.h"
#include "paddle/phi/capi/include/c_meta_tensor.h"
#include "paddle/phi/capi/include/c_place.h"
#include "paddle/phi/capi/include/c_scalar.h"
#include "paddle/phi/capi/include/c_tensor.h"
#include "paddle/phi/capi/include/data_type.h"
#include "paddle/phi/capi/include/kernel_registry.h"
#endif
因此paddle为custom device提供的头文件全部都在paddle/phi/capi
文件夹中,由于只include了capi文件夹中的头文件,因此AddKernel函数参数中的phi::DenseTensor
数据类型实际上并不是paddle/phi/core/dense_tensor.h
中定义的phi::DenseTensor
,而是paddle/phi/capi/include/kernel_registry.h
中如下定义的数据类型:
cpp
namespace phi {
.....
using DenseTensor = capi::DenseTensor;
.....
}
因此,custom_cpu中算子使用的phi::DenseTensor
实际上是paddle/phi/capi/include/wrapper_base.h
中定义的phi::capi::DenseTensor
,同理,custom_cpu中算子使用的phi::Context
实际上是paddle/phi/capi/include/wrapper_base.h
中定义的phi::capi::DeviceContext
。
2 phi::capi::DenseTensor讲解
该class的定义位于paddle/phi/capi/include/wrapper_base.h
,定义如下:
cpp
template <typename T>
class WrapperBase {
public:
explicit WrapperBase(T* ptr, bool own = false) : data_(ptr), own_(own) {}
inline T* raw_data() const { return data_; }
inline bool own_data() const { return own_; }
inline void reset(const T* ptr) { data_ = ptr; }
private:
T* data_;
bool own_;
};
// PD_Tensor定义于paddle/phi/capi/include/c_tensor.h
// 定义内容为:typedef struct PD_Tensor PD_Tensor;
// 即PD_Tensor为一个空结构体
class DenseTensor : public WrapperBase<PD_Tensor> {
//继承WrapperBase<PD_Tensor>,因此DenseTensor有一个PD_Tensor类型的指针作为私有成员变量
public:
DenseTensor() : WrapperBase(PD_NewTensor(), true) {}
explicit DenseTensor(PD_Tensor* tensor) : WrapperBase(tensor) {}
~DenseTensor() {
if (own_data()) {
PD_DeleteTensor(raw_data());
}
}
size_t offset() const {
C_Status status;
//PD_TensorGetOffset函数传入的第一个参数是PD_Tensor类型的指针
auto offset = PD_TensorGetOffset(raw_data(), &status);
PD_CHECK_STATUS(status);
return offset;
}
............................................
}
//以下是定义于paddle/phi/capi/lib/c_tensor.cc的PD_TensorGetOffset函数
size_t PD_TensorGetOffset(const PD_Tensor* tensor, PD_Status* status) {
if (status) {
if (!tensor) {
*status = C_FAILED;
return 0;
}
*status = C_SUCCESS;
}
auto cc_tensor = reinterpret_cast<const phi::DenseTensor*>(tensor); 将指针转为phi::DenseTensor类型并调用其方法
return cc_tensor->offset();
}