Paadle Custom device 代码阅读记录

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();
}
相关推荐
聚客AI11 小时前
🙋‍♀️Transformer训练与推理全流程:从输入处理到输出生成
人工智能·算法·llm
大怪v14 小时前
前端:人工智能?我也会啊!来个花活,😎😎😎“自动驾驶”整起!
前端·javascript·算法
惯导马工16 小时前
【论文导读】ORB-SLAM3:An Accurate Open-Source Library for Visual, Visual-Inertial and
深度学习·算法
骑自行车的码农17 小时前
【React用到的一些算法】游标和栈
算法·react.js
博笙困了18 小时前
AcWing学习——双指针算法
c++·算法
moonlifesudo18 小时前
322:零钱兑换(三种方法)
算法
NAGNIP1 天前
大模型框架性能优化策略:延迟、吞吐量与成本权衡
算法
美团技术团队2 天前
LongCat-Flash:如何使用 SGLang 部署美团 Agentic 模型
人工智能·算法
Fanxt_Ja2 天前
【LeetCode】算法详解#15 ---环形链表II
数据结构·算法·leetcode·链表
侃侃_天下2 天前
最终的信号类
开发语言·c++·算法