找回密码
立即注册
搜索
热搜: Java Python Linux Go
发回帖 发新帖

1459

积分

0

好友

187

主题
发表于 2026-2-14 03:45:01 | 查看: 32| 回复: 0

本文详细阐述一套基于 C++17 标准库、完全脱离 PyTorch 的 Tensor 运算分发架构。该架构通过 全局算子注册表类型 ID 隔离构造时分配内存 等设计,实现了 零侵入添加新数据类型 的核心目标。核心模块与扩展插件分处不同目录,扩展方只需包含核心头文件、调用注册函数,无需修改任何核心代码 即可为自定义类型注册完整的运算支持。全文从类型系统、Tensor 实现、注册机制、内置实现到扩展示例逐一剖析,完整呈现一个可工业级复用的静态分发原型。

整体架构概览

下图描述了系统的核心模块及其交互关系。所有模块均通过全局注册表解耦,核心与扩展无直接依赖。

Tensor系统架构流程图

  • Tensor:仅持有类型 ID 与数据指针(shared_ptr<void>),不包含任何算子逻辑。
  • 类型大小表:存储每种 ScalarType 对应的字节大小,Tensor 构造时据此分配内存。
  • 算子表:以 "op_lhs_rhs" 为键,存储二元 Kernel 函数,由 add 查询调用。
  • Kernel 实现:内置与扩展分别提供,仅通过注册表与分发函数耦合。

所有扩展操作均通过注册接口完成,核心组件无需预知任何自定义类型。

核心组件详解

3.1 类型标识系统

using ScalarType = std::uint32_t;
constexpr ScalarType kFloat = 0;
constexpr ScalarType kInt   = 1;
// 自定义类型从较大数值开始,如 1000

类型 ID 分配策略如下表所示,内置与自定义区间物理隔离,永不重叠

类型 可供分配的 ID
内置类型 0‑999
自定义类型 ≥1000
  • ScalarType 定义为 uint32_t,完全抛弃 enum class,便于扩展方自由定义新常量。
  • 内置类型 ID 从 0 开始递增,集中定义在 core/core.hpp 中。
  • 自定义类型 ID 从 1000 或其他大数开始,定义在 extension/extension.hpp 中,与内置 ID 绝对不重叠。
  • 此隔离策略是零侵入的基石:核心代码从不引用扩展的 ID,扩展也不会影响内置 ID 序列。

3.2 Tensor 类设计

class Tensor {
public:
    explicit Tensor(ScalarType type);          // 唯一公开构造函数
    template <typename T> void fill(const T& value);
    void* data() const noexcept;
    ScalarType scalar_type() const noexcept;
private:
    ScalarType type_;
    std::shared_ptr<void> data_;
};
  • 构造函数仅接受类型 ID,不接收任何数据指针,杜绝外部非法构造
  • 构造时根据 type_ 查询全局 type_size 表,调用 operator new 分配原始内存,存入 shared_ptr<void>,删除器为 operator delete
  • fill 模板:将用户提供的值直接拷贝至已分配内存。要求 sizeof(T) 必须等于注册的类型大小,否则抛出异常。
  • data() 返回裸指针,仅用于 Kernel 内部读写,所有权仍由 shared_ptr 管理

这种设计将“内存分配”与“数值赋值”彻底分离:分配发生在构造时,赋值通过 fill 完成,Kernel 只需构造结果 Tensor 并直接写入结果,无需关心内存管理细节。

3.3 全局注册表

extern std::unordered_map<std::string, KernelFunc> kernels;
extern std::unordered_map<ScalarType, size_t> type_size;

void register_kernel(const std::string& op, ScalarType lhs, ScalarType rhs, KernelFunc func);
void register_type_size(ScalarType type, size_t size);

注册表内部结构及注册流程如下图所示:

系统注册流程示意图

  • 算子表:键为 "add_0_0" 格式的字符串,值是可调用对象 Tensor(const Tensor&, const Tensor&)
  • 类型大小表:键为 ScalarType,值为 sizeof(T)Tensor 构造时必须依赖此表
  • 两表均为 全局静态变量,定义在 core.cpp 中,仅通过注册接口暴露。

注册接口不依赖任何具体类型,扩展方可随时调用,将自定义 Kernel 和类型大小插入全局表。

3.4 算子分发函数

add 为例:

Tensor add(const Tensor& a, const Tensor& b){
    std::string key = "add_" + std::to_string(a.scalar_type()) + "_"
                     + std::to_string(b.scalar_type());
    auto it = kernels.find(key);
    if (it != kernels.end()) return it->second(a, b);
    throw std::runtime_error("No kernel");
}
  • 运行时根据左右操作数的类型 ID 动态拼接键名,完全避免静态分支
  • 查表成功则调用对应 Kernel,失败则抛异常。
  • 新增类型组合无需修改此函数,只需注册对应的 Kernel 即可。

内置类型实现(core 目录)

4.1 类型注册与内存规格

void register_builtin(){
    register_type_size(kFloat, sizeof(float));
    register_type_size(kInt,   sizeof(int));
    register_kernel("add", kFloat, kFloat, float_add_float);
    register_kernel("add", kInt,   kInt,   int_add_int);
}
  • 内置类型的大小与 Kernel 必须在核心初始化时注册,否则无法构造 Tensor 或调用 add
  • 类型大小注册 必须在任何 Tensor 构造之前完成,由用户显式调用 register_builtin() 保证。

4.2 内置 Kernel 实现

Tensor float_add_float(const Tensor& a, const Tensor& b){
    Tensor result(kFloat);               // 已分配 float 内存
    float* pa = static_cast<float*>(a.data());
    float* pb = static_cast<float*>(b.data());
    float* pres = static_cast<float*>(result.data());
    *pres = *pa + *pb;
    return result;
}

Tensor int_add_int(const Tensor& a, const Tensor& b){
    Tensor result(kInt);                // 已分配 int 内存
    int* pa = static_cast<int*>(a.data());
    int* pb = static_cast<int*>(b.data());
    int* pres = static_cast<int*>(result.data());
    *pres = *pa + *pb;
    return result;
}
  • Kernel 不再承担内存分配责任:结果 Tensor 构造时即分配好内存。
  • 直接通过 data() 获取裸指针并写入计算结果,代码简洁且异常安全shared_ptr 自动管理)。

这种模式使 Kernel 逻辑聚焦于数值计算,与内存分配解耦,极大降低了 Kernel 编写错误的风险。

4.3 核心代码全貌(core.hpp / core.cpp)

为使读者能够完整运行,下面展示 核心模块的全部源代码

core/core.hpp

#ifndef CORE_HPP
#define CORE_HPP

#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

namespace at {

using ScalarType = std::uint32_t;
constexpr ScalarType kFloat = 0;
constexpr ScalarType kInt   = 1;

class Tensor {
public:
    explicit Tensor(ScalarType type);

    Tensor(Tensor&&) noexcept = default;
    Tensor& operator=(Tensor&&) noexcept = default;
    Tensor(const Tensor&) = delete;
    Tensor& operator=(const Tensor&) = delete;
    ~Tensor() = default;

    ScalarType scalar_type() const noexcept{ return type_; }
    void* data() const noexcept{ return data_.get(); }

    template <typename T>
    void fill(const T& value);

private:
    ScalarType type_;
    std::shared_ptr<void> data_;
    std::vector<int64_t> shape_;   // 占位,不实现广播
};

using KernelFunc = std::function<Tensor(const Tensor&, const Tensor&)>;

extern std::unordered_map<std::string, KernelFunc> kernels;
extern std::unordered_map<ScalarType, size_t> type_size;

void register_kernel(const std::string& op, ScalarType lhs, ScalarType rhs,
                     KernelFunc func);
void register_type_size(ScalarType type, size_t size);

Tensor add(const Tensor& a, const Tensor& b);

Tensor float_add_float(const Tensor& a, const Tensor& b);
Tensor int_add_int(const Tensor& a, const Tensor& b);

void register_builtin();

} // namespace at

#endif // CORE_HPP

core/core.cpp

#include "core/core.hpp"
#include <cstring>
#include <stdexcept>

namespace at {

std::unordered_map<std::string, KernelFunc> kernels;
std::unordered_map<ScalarType, size_t> type_size;

Tensor::Tensor(ScalarType type) : type_(type) {
    auto it = type_size.find(type);
    if (it == type_size.end()) {
        throw std::runtime_error("Cannot create Tensor: dtype size not registered");
    }
    data_ = std::shared_ptr<void>(operator new(it->second),
                                 [](void* p) { operator delete(p); });
}

template <typename T>
void Tensor::fill(const T& value){
    auto it = type_size.find(type_);
    if (it == type_size.end())
        throw std::runtime_error("Unknown dtype, size not registered");
    if (sizeof(T) != it->second)
        throw std::runtime_error("sizeof(T) does not match registered dtype size");
    *static_cast<T*>(data_.get()) = value;
}

template void Tensor::fill<float>(const float&);
template void Tensor::fill<int>(const int&);

void register_kernel(const std::string& op, ScalarType lhs, ScalarType rhs,
                     KernelFunc func){
    std::string key = op + "_" + std::to_string(lhs) + "_" + std::to_string(rhs);
    kernels[key] = std::move(func);
}

void register_type_size(ScalarType type, size_t size){
    type_size[type] = size;
}

Tensor add(const Tensor& a, const Tensor& b){
    std::string key = "add_" + std::to_string(a.scalar_type()) + "_"
                     + std::to_string(b.scalar_type());
    auto it = kernels.find(key);
    if (it != kernels.end()) return it->second(a, b);
    throw std::runtime_error("No kernel for add with given types");
}

Tensor float_add_float(const Tensor& a, const Tensor& b){
    Tensor result(kFloat);
    *static_cast<float*>(result.data()) = *static_cast<float*>(a.data()) +
                                          *static_cast<float*>(b.data());
    return result;
}

Tensor int_add_int(const Tensor& a, const Tensor& b){
    Tensor result(kInt);
    *static_cast<int*>(result.data()) = *static_cast<int*>(a.data()) +
                                        *static_cast<int*>(b.data());
    return result;
}

void register_builtin(){
    register_type_size(kFloat, sizeof(float));
    register_type_size(kInt,   sizeof(int));
    register_kernel("add", kFloat, kFloat, float_add_float);
    register_kernel("add", kInt,   kInt,   int_add_int);
}

} // namespace at

零侵入扩展机制(extension 目录)

5.1 目录结构隔离

项目采用 物理目录完全分离 的结构,核心与扩展无任何编译时依赖:

project/
├── core/
│   ├── core.hpp
│   └── core.cpp
├── extension/
│   ├── extension.hpp
│   └── extension.cpp
└── main.cpp

项目文件结构依赖图

  • 物理目录完全隔离,核心代码不包含任何扩展头文件,扩展仅通过 #include "core/core.hpp" 依赖核心。
  • 编译时通过 -I. 包含根目录,核心与扩展可分别编译成静态库后再链接,实现二进制级别的隔离。

5.2 自定义类型定义

namespace at {
    constexpr ScalarType kMyType = 1000;
    struct MyTypeData {
        int8_t value;
        float scale;
    };
}
  • 自定义类型 ID 选择不与内置冲突的大数,保证唯一性。
  • 自定义数据结构完全自由,不受核心任何约束(只需与注册的大小一致)。

5.3 自定义 Kernel 实现

Tensor mytype_add_mytype(const Tensor& a, const Tensor& b){
    Tensor result(kMyType);
    auto* da = static_cast<MyTypeData*>(a.data());
    auto* db = static_cast<MyTypeData*>(b.data());
    float fa = da->value * da->scale;
    float fb = db->value * db->scale;
    float fsum = fa + fb;
    float new_scale = da->scale + db->scale;
    int8_t quant = static_cast<int8_t>(std::round(fsum / new_scale));
    *static_cast<MyTypeData*>(result.data()) = MyTypeData{quant, new_scale};
    return result;
}
  • 与内置 Kernel 遵循完全相同的模式:构造结果 Tensor,写入数据,返回。
  • 无需关心内存分配/释放,所有资源管理由 Tensor 的 shared_ptr 统一处理

5.4 扩展注册流程

void register_extension(){
    register_type_size(kMyType, sizeof(MyTypeData));
    register_kernel("add", kMyType, kMyType, mytype_add_mytype);
}
  • 扩展模块通过一个独立的初始化函数完成所有注册,该函数由用户在主程序中调用。
  • 整个注册过程 零修改核心代码,核心不包含任何与 MyType 相关的预处理宏、枚举项或条件编译。

这就是“零侵入”的完整定义:扩展方只需要编写自己的类型与 Kernel,调用核心提供的注册 API,即可无缝融入系统。

5.5 扩展代码全貌(extension.hpp / extension.cpp)

extension/extension.hpp

#ifndef EXTENSION_HPP
#define EXTENSION_HPP

#include "core/core.hpp"

namespace at {

constexpr ScalarType kMyType = 1000;

struct MyTypeData {
    int8_t value;
    float scale;
};

Tensor mytype_add_mytype(const Tensor& a, const Tensor& b);
void register_extension();

} // namespace at

#endif // EXTENSION_HPP

extension/extension.cpp

#include "extension/extension.hpp"
#include <cmath>

namespace at {

Tensor mytype_add_mytype(const Tensor& a, const Tensor& b){
    Tensor result(kMyType);
    auto* da = static_cast<MyTypeData*>(a.data());
    auto* db = static_cast<MyTypeData*>(b.data());
    float fa = da->value * da->scale;
    float fb = db->value * db->scale;
    float fsum = fa + fb;
    float new_scale = da->scale + db->scale;
    int8_t quant = static_cast<int8_t>(std::round(fsum / new_scale));
    *static_cast<MyTypeData*>(result.data()) = MyTypeData{quant, new_scale};
    return result;
}

void register_extension(){
    register_type_size(kMyType, sizeof(MyTypeData));
    register_kernel("add", kMyType, kMyType, mytype_add_mytype);
}

} // namespace at

编译与运行演示

6.1 主程序(main.cpp)

#include "core/core.hpp"
#include "extension/extension.hpp"
#include <iostream>

int main(){
    using namespace at;

    register_builtin();
    register_extension();

    // ----- Float -----
    Tensor t1(kFloat);
    t1.fill<float>(1.2f);
    Tensor t2(kFloat);
    t2.fill<float>(3.4f);
    Tensor t3 = add(t1, t2);
    std::cout << "Float add: " << *(float*)t3.data() << std::endl;

    // ----- Int -----
    Tensor ti1(kInt);
    ti1.fill<int>(5);
    Tensor ti2(kInt);
    ti2.fill<int>(7);
    Tensor ti3 = add(ti1, ti2);
    std::cout << "Int add: " << *(int*)ti3.data() << std::endl;

    // ----- MyType -----
    Tensor q1(kMyType);
    q1.fill<MyTypeData>({10, 0.1f});   // 1.0
    Tensor q2(kMyType);
    q2.fill<MyTypeData>({20, 0.2f});   // 4.0
    Tensor q3 = add(q1, q2);
    auto* res = static_cast<MyTypeData*>(q3.data());
    std::cout << "MyType add: quantized=" << (int)res->value
              << ", scale=" << res->scale
              << ", real=" << (res->value * res->scale) << std::endl;

    return 0;
}

6.2 编译命令与输出

# 从项目根目录执行
g++ -std=c++17 -I. -o demo core/core.cpp extension/extension.cpp main.cpp
./demo

预期输出

Float add: 4.6
Int add: 12
MyType add: quantized=15, scale=0.3, real=4.5
  • 内置类型与自定义类型使用完全一致的 APITensor t(kMyType); + t.fill<MyTypeData>(...) + add(t1, t2)
  • 用户代码 无需区分内置与扩展,系统在运行时自动分派到正确的 Kernel。

对比分析:传统扩展 vs 零侵入扩展

维度 传统扩展(硬编码) 零侵入扩展(本架构)
修改核心枚举 必须向 ScalarType 添加新枚举值 无需修改,自定义 ID 独立定义
修改分发函数 必须在 add 中添加 switch 分支 无需修改add 纯查表
修改核心 Kernel 表 需将新 Kernel 硬链接到全局数组或宏列表中 动态注册,运行时调用 register_kernel
内存管理适配 需修改 Tensor 构造逻辑或提供新工厂 Tensor 构造已通用,仅需注册类型大小
编译依赖 扩展代码通常侵入核心头文件,导致全量重编译 核心与扩展分离编译,扩展更新无需重编核心
二进制兼容性 核心库必须为每种新类型重新发布 核心库可保持稳定,扩展以插件形式动态/静态链接

此表清晰揭示零侵入架构的本质核心系统对扩展类型完全无知,所有扩展信息均通过注册表在运行时注入。这种设计模式上的解耦,使得系统具备了极高的可维护性和可扩展性。想了解更多关于类似设计模式和系统设计思想的讨论,可以访问 云栈社区 的开发者论坛。

总结

本文完整呈现了一套纯 C++17 的 Tensor 分发系统,其核心贡献在于:

  1. 类型 ID 隔离策略:内置与自定义 ID 分属不同数值区间,从根本上避免了枚举冲突。
  2. 内存分配与数值赋值分离:Tensor 构造时统一分配原始内存,fill 模板负责类型安全赋值,Kernel 仅关注计算。
  3. 全局注册表驱动的分派:算子表与类型大小表均为运行时动态结构,任何新类型只需调用注册接口即可接入。
  4. 物理目录隔离:核心与扩展分处不同文件夹,不产生任何编译期依赖,实现二进制级别的可扩展性。

该架构已在原型中成功支持 内置浮点/整型自定义量化 int8 类型 的无缝混合运算,所有扩展代码均未触及核心模块。这套设计可直接移植到需要支持用户定义数据类型的 C++ 数值计算库中,也可作为深度学习框架自定义后端的底层参考实现,其背后的思想对于理解现代软件库的 内存管理 和扩展性设计颇具启发性。




上一篇:字节ByteRobust解析:保障大规模LLM训练的自动容错系统
下一篇:白帽研究会被“误伤”吗?解读《网络犯罪防治法》征求意见稿第24、25条对安全从业者的影响
您需要登录后才可以回帖 登录 | 立即注册

手机版|小黑屋|网站地图|云栈社区 ( 苏ICP备2022046150号-2 )

GMT+8, 2026-2-23 15:21 , Processed in 0.354889 second(s), 41 queries , Gzip On.

Powered by Discuz! X3.5

© 2025-2026 云栈社区.

快速回复 返回顶部 返回列表