diff --git a/Makefile b/Makefile index 35ef7efd..f05cb98e 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY : build clean format install-python test-cpp test-onnx +.PHONY : build clean format install-python test-cpp test-onnx submodules TYPE ?= Release TEST ?= ON @@ -6,10 +6,15 @@ TEST ?= ON CMAKE_OPT = -DCMAKE_BUILD_TYPE=$(TYPE) CMAKE_OPT += -DBUILD_TEST=$(TEST) -build: +# 注意：build 必须是第一个目标，否则 `make` 的默认目标会变成 submodules +build: submodules mkdir -p build/$(TYPE) cd build/$(TYPE) && cmake $(CMAKE_OPT) ../.. && make -j8 +# 初始化 git submodule（googletest 等第三方依赖，clone 后首次构建必须执行） +submodules: + git submodule update --init --recursive + clean: rm -rf build diff --git "a/docs/\344\275\234\344\270\232\345\256\214\346\210\220\350\257\264\346\230\216.md" "b/docs/\344\275\234\344\270\232\345\256\214\346\210\220\350\257\264\346\230\216.md" new file mode 100644 index 00000000..051cae8b --- /dev/null +++ "b/docs/\344\275\234\344\270\232\345\256\214\346\210\220\350\257\264\346\230\216.md" @@ -0,0 +1,1026 @@ +# TinyInfiniTensor 训练营作业完成说明 + +## 完成情况 + +本文档对应 [训练营作业介绍](./训练营作业介绍.md) 中的 8 道作业题，**已全部完成并通过本地自测**： + +| # | 作业 | 难度 | 对应测例 | 状态 | +| --- | --- | --- | --- | --- | +| 一 | 内存分配器（allocator + dataMalloc） | ⭐⭐⭐⭐ | test_allocator、test_nativecpu_concat、test_nativecpu_elementwise、test_nativecpu_transpose | ✅ 已完成 | +| 二 | transpose 算子形状推导 | ⭐ | test_transpose、test_nativecpu_transpose | ✅ 已完成 | +| 三 | clip 算子形状推导 | ⭐ | test_clip | ✅ 已完成 | +| 四 | cast 形状推导 & 数据类型推导 | ⭐⭐ | test_cast | ✅ 已完成 | +| 五 | concat 算子形状推导 | ⭐⭐ | test_concat、test_nativecpu_concat | ✅ 已完成 | +| 六 | 双向广播 | ⭐⭐⭐ | test_element_wise、test_nativecpu_elementwise、test_matmul | ✅ 已完成 | +| 七 | 矩阵乘形状推导 | ⭐⭐⭐ | test_matmul | ✅ 已完成 | +| 八 | 简单图优化规则实现 | ⭐⭐⭐⭐ | test_graph | ✅ 已完成 | + +**测试结果**：`make test-cpp` 11/11 全部通过（每个测例 10 分）。 + +## 文档定位 + +本文档是作业的**完成记录与实现说明**：记录了每个作业的实现思路、关键原理、 +完整参考实现（即当前仓库中的代码）以及开发过程中踩过的坑。自测命令： + +```bash +make build # 编译 +make test-cpp # 运行全部 11 个测例 +``` + +以下各节按作业顺序记录实现说明（顺序：二 → 三 → 四 → 五 → 六 → 七 → 一 → 八）。 + +--- + +## 0. 项目速览与常用概念 + +### 0.1 目录结构与构建 + +```text +include/ 头文件（声明） + core/ tensor / operator / graph / allocator / runtime / blob ... + operators/ transpose / unary / concat / matmul / element_wise 的算子类 + utils/ 工具函数（广播、异常等） +src/ + core/ core 组件的实现 + operators/ 各算子实现（形状推导 inferShape 在这里） + kernels/cpu/ CPU kernel（真正做计算的代码，作业不涉及） + utils/ 工具函数实现（infer_broadcast 在这里） +test/ + core/ test_allocator、test_graph + operators/ test_transpose、test_clip、test_cast、test_concat、test_element_wise、test_matmul + kernels/nativecpu/ test_nativecpu_concat、test_nativecpu_elementwise、test_nativecpu_transpose +docs/ 项目文档 +``` + +构建与自测（进入项目根目录）： + +```bash +make build # 默认 Release 模式构建（含测试） +make test-cpp # 运行全部 11 个测例，全部通过即完成作业 +``` + +### 0.2 核心对象模型 + +| 类型 | 说明 | +| --- | --- | +| `TensorObj` | 一张量，包含形状 `Shape`（`vector`）、数据类型 `DataType`、数据内存 `Blob`。`getBytes()` 返回占用字节数（`size() * dtype.getSize()`） | +| `OperatorObj` | 一个算子，持有输入 `inputs`、输出 `outputs`，以及前驱 `predecessors`/后继 `successors` | +| `GraphObj` | 计算图，管理 `tensors`、`ops`、`allocator` | +| `BlobObj` | 封装一块真实内存指针，`tensor->setDataBlob(blob)` 把内存绑定给 tensor | +| `Ref` / `WRef` | 分别是 `shared_ptr` / `weak_ptr` 的别名。tensor 与算子之间用 `weak_ptr` 互相引用，避免循环引用 | + +图中有两组双向关系，**写代码维护关系时最容易在这里出错**： + +- **tensor 侧**：`source`（谁产生了我，`weak_ptr`）与 `targets`（谁消费我，`weak_ptr` 列表）。 +- **op 侧**：`predecessors` / `successors`。 + +这两个视图必须同步维护。例如 `addOperatorAndConnect` 中会同时维护这两侧。 + +### 0.3 形状推导流程（理解作业二~七的关键） + +以 `g->addOp(i, nullptr, perm)` 为例，创建算子时会走： + +```cpp +OperatorObj::checkValid(graph) // src/core/operator.cc:43 + ├─ inferShape() // 调用各算子的形状推导 → 得到输出形状 + ├─ inferDataType() // 得到输出数据类型 + └─ outputs[i] = graph->addTensor(...) // 按推导出的形状/类型创建输出 tensor +``` + +所以作业二~七的任务，本质是**让 `inferShape` 返回正确形状**（部分算子还要重写 `inferDataType`）。 +返回类型是 `optional>`，每个元素对应一个输出的形状；返回 `std::nullopt` 表示推导失败。 + +注意：`checkValid` 中已经调用了 `inferShape`，所以 `g->addOp(...)` 返回的 +算子对象的输出 tensor 就已经带上了正确的形状，测例都是通过 +`op->getOutput()->getDims()` 来断言的。 + +### 0.4 guid 与 Fuid + +- `guid`：每个 `Object`（tensor、op、graph）全局唯一的 ID，从 1 开始递增。 +- `fuid`：tensor 的"家族 ID"，克隆的 tensor 共享同一个 fuid。 +- `test_graph` 中 `EXPECT_EQ(op->getInputs(0)->getGuid(), 2)` 就是利用 + guid 从 1 开始（graph 占 1，后续创建的 tensor 依次为 2、3、…）来验证重连结果。 + +--- + +## 作业一：内存分配器（⭐⭐⭐⭐） + +**对应测例**：`test_allocator`、`test_nativecpu_concat`、`test_nativecpu_elementwise`、`test_nativecpu_transpose` + +**涉及文件**：`include/core/allocator.h`、`src/core/allocator.cc`、`src/core/graph.cc`（dataMalloc） + +### 1.1 设计要点：这是一个"模拟"分配器 + +先理解 `Allocator` 的总体设计（`include/core/allocator.h` + `src/core/allocator.cc`）： + +- `alloc(size)` / `free(addr, size)` 都带有 `IT_ASSERT(this->ptr == nullptr)`， + 也就是说 **alloc/free 阶段完全不做真实内存操作**，只是"纸上谈兵"地计算偏移量。 +- 真实内存只在 `getPtr()` 里分配**一次**（大小等于 `peak`，即规划出的最大内存）， + 之后所有 tensor 都通过 `基址 + 偏移量` 的方式共享这块内存。 +- `used` 是当前"模拟"状态下被占用的字节数；`peak` 是历史上 `used` 达到的最大值 + （即真正需要分配的内存大小）；`getPtr()` 就是按 `peak` 分配。 + +因此我们的任务就是实现两件事： + +1. 用 `freeBlocks` 记录空闲块，实现 `alloc`（找到合适空闲块或从尾部新分配）与 + `free`（回收并合并相邻空闲块）； +2. 在 `GraphObj::dataMalloc` 里**先规划**好每个 tensor 的偏移量，**最后**一次性 + `getPtr()` 绑定真实内存。 + +### 1.2 step1：声明空闲块数据结构 + +`include/core/allocator.h` 的 private 区： + +```cpp +// =================================== 作业 =================================== +// std::map：key 为 free block 的起始地址偏移量， +// value 为 block 的大小。map 会按 key（地址）自动升序排列， +// 这样既能快速找到满足大小的空闲块，也能方便地与相邻块做合并。 +std::map freeBlocks; +// =================================== 作业 =================================== +``` + +**为什么用 `std::map`？** + +- 天然按地址升序排列，可以用 first-fit 策略从低地址开始扫描； +- `std::prev(it)` / `std::next(it)` 可以方便地访问相邻块做合并； +- 插入/删除是 O(log n)，对学习场景足够。 + +### 1.3 step2：实现 `alloc` + +```cpp +size_t Allocator::alloc(size_t size) +{ + IT_ASSERT(this->ptr == nullptr); + // pad the size to the multiple of alignment + size = this->getAlignedSize(size); + + // =================================== 作业 =================================== + // 采用 first-fit 策略：从低地址到高地址扫描 free block， + // 找到第一个大小不小于 size 的空闲块就进行分配。 + for (auto it = freeBlocks.begin(); it != freeBlocks.end(); ++it) + { + if (it->second >= size) // 找到足够大的空闲块 + { + size_t addr = it->first; // 该块的起始地址就是返回的偏移量 + // 空闲块比需要的更大：把剩余部分重新放回 free list + if (it->second > size) + { + freeBlocks[addr + size] = it->second - size; + } + freeBlocks.erase(it); // 这个块被整块用掉（或切掉一部分） + used += size; // 更新模拟使用量 + peak = std::max(peak, used); // 更新历史峰值 + return addr; + } + } + // 没有合适的空闲块，就在已用内存的末尾追加新的内存。 + // 特例：如果地址最大的空闲块一直延伸到当前内存末尾（peak）， + // 可以从该块的起始地址继续向外扩展，从而复用之前被释放的"末尾块"。 + if (!freeBlocks.empty()) + { + auto last = std::prev(freeBlocks.end()); + if (last->first + last->second == peak) + { + size_t addr = last->first; + peak = addr + size; // 内存末尾向后扩展 + used += size; + freeBlocks.erase(last); + return addr; + } + } + size_t addr = peak; // 普通情况：直接追加到内存末尾 + peak += size; + used += size; + return addr; +} +``` + +**逐段解释：** + +1. **first-fit 扫描**：从低地址开始找第一个 `second >= size` 的空闲块。找到了就把 + `[addr, addr+size)` 分配出去；如果空闲块更大，把 `[addr+size, addr+second)` 的 + 剩余部分作为新的空闲块塞回 map。这里 `freeBlocks[addr + size] = ...` 利用 + map 的自动排序，新块会落在正确的位置。 + +2. **末尾块扩展（容易漏掉的细节）**：`test_allocator` 的 + `testAllocWithEndFreeBlock` 故意设计了这样一个场景： + - 依次分配 a、b、c 三块内存（c 在末尾，`[96,144)`）； + - free c，此时 free list 里只有 `{[96,144):48}` 这一个块； + - 再分配 d（96 字节，是 c 的两倍大）。 + - 测例期望 `offsetD == offsetC`，即 d 复用 c 的位置。 + + 如果只做 first-fit，48 字节的空闲块装不下 96 字节，d 会被放到 144 处， + 测例失败。正确的做法是：**当最后一个空闲块正好延伸到内存末尾（`peak`）时**， + 可以从它的起始地址继续向外扩展。`96+48 == 144 == peak`，因此 d 从 96 开始， + 扩展到 `96+96=192`，既复用了 c 的位置，又不会与 b 重叠。 + +3. **追加分配**：既没有合适的空闲块、也没有可扩展的末尾块时，在 `peak` 处追加。 + 注意这里返回的 `addr = peak`（旧的 peak），然后把 peak 向后推。 + +### 1.4 step2：实现 `free` + +```cpp +void Allocator::free(size_t addr, size_t size) +{ + IT_ASSERT(this->ptr == nullptr); + size = getAlignedSize(size); + + // =================================== 作业 =================================== + used -= size; + auto it = freeBlocks.emplace(addr, size).first; + + // 合并前一块：前一块的结束地址 == 当前块的起始地址 + if (it != freeBlocks.begin()) + { + auto prev = std::prev(it); + if (prev->first + prev->second == addr) + { + prev->second += it->second; // 把当前块并入前一块 + freeBlocks.erase(it); + it = prev; + } + } + // 合并后一块：当前块的结束地址 == 后一块的起始地址 + auto next = std::next(it); + if (next != freeBlocks.end() && it->first + it->second == next->first) + { + it->second += next->second; // 把后一块并入当前块 + freeBlocks.erase(next); + } + // =================================== 作业 =================================== +} +``` + +**解释：** + +- `free` 的核心是"回收 + 合并"。如果不合并，连续 free 两个相邻块会产生两个 + 碎片化的空闲块，后续的大块分配可能因为凑不齐而失败，或者导致不必要的内存增长。 +- 合并分两步：**先向前合并，再向后合并**（顺序不能反，向前合并后 `it` 指向 + 合并后的块，再用它向后合并，否则可能漏掉第二次合并）。 +- 合并条件是**地址恰好相邻**：`前一块的 first + second == 当前块的 addr`。 + +### 1.5 step3：`GraphObj::dataMalloc`（计算图内存分配） + +```cpp +void GraphObj::dataMalloc() +{ + // topological sorting first + IT_ASSERT(topo_sort() == true); + + // =================================== 作业 =================================== + std::map refCount; // tensor -> 剩余消费者个数 + std::map offsetMap; // tensor -> 规划到的偏移量 + std::unordered_set allocated; // 已规划过内存的 tensor + for (auto &t : tensors) + refCount[t] = t->getTargets().size(); + + // 阶段一：模拟分配，规划每个 tensor 的偏移量 + for (auto &op : ops) + { + // 为算子的输入分配内存（每个 tensor 只规划一次） + for (auto &t : op->getInputs()) + { + if (allocated.find(t) == allocated.end()) + { + offsetMap[t] = allocator.alloc(t->getBytes()); + allocated.insert(t); + } + } + // 为算子的输出分配内存（先于输入回收，避免与输入重叠） + for (auto &t : op->getOutputs()) + { + if (allocated.find(t) == allocated.end()) + { + offsetMap[t] = allocator.alloc(t->getBytes()); + allocated.insert(t); + } + } + // 本算子的输入在本算子执行完后不再被它使用： + // 引用计数减一，归零的输入（最后一个消费者已完成）立即回收， + // 供后面的算子复用，从而降低峰值内存。 + for (auto &t : op->getInputs()) + { + if (--refCount[t] == 0) + allocator.free(offsetMap[t], t->getBytes()); + } + } + + // 阶段二：真正分配内存，并按偏移量给每个 tensor 绑定 Blob + auto base = allocator.getPtr(); + for (auto &[t, offset] : offsetMap) + t->setDataBlob(make_ref(runtime, (char *)base + offset)); + // =================================== 作业 =================================== + + allocator.info(); +} +``` + +**逐段解释（这一段是作业一里最容易出问题的部分）：** + +1. **为什么必须分两个阶段？** + 因为 allocator 的 `alloc/free` 带断言 `IT_ASSERT(this->ptr == nullptr)`， + `getPtr()` 只能调用一次且必须在所有模拟分配**之后**。如果边 alloc 边 + `getPtr()` 绑 Blob，第二次 alloc 就直接断言失败。 + +2. **`refCount` 的用途（内存复用）**：每个 tensor 初始化时记录它还有多少个消费者 + （即 `targets` 的个数）。当拓扑序遍历到某个 op 时，它的输入被"消费"一次， + 计数减一；减到 0 说明该 tensor 之后再也不会被读，内存可以 `free` 掉， + 让后面的算子复用。例如一个 `i1 -> op1 -> o1 -> op2 -> o2` 的链条， + `i1` 在 op1 规划完就被回收，`o2` 就可以复用 `i1` 的内存，peak 内存近似减半。 + +3. **输出必须先于输入回收来分配（关键顺序）**：在一个算子内部，如果先回收输入再 + 分配输出，输出会复用**本算子输入**的内存。而 kernel 实际执行时是 + **边读输入边写输出**的（见 `src/kernels/cpu/concat.cc` 的 + `outPtr[oOffset] = inPtr[iOffset]`），输入与输出重叠就会算错。 + 因此代码把「分配输出」放在「回收输入」之前。 + +4. **为什么要 `allocated` 集合**：同一个 tensor 可能被多个算子消费（多输入分支）， + 如果不做去重，第二个算子会给同一个 tensor 再 alloc 一块内存，浪费内存且 + `offsetMap` 会被覆盖。`allocated` 保证每个 tensor 只规划一次。 + +5. **绑定 Blob**：`allocator.getPtr()` 返回真实内存基址，加上偏移量后构造 + `BlobObj`，通过 `tensor->setDataBlob(...)` 绑定。之后测试里的 + `t->setData(...)`（写数据）与 `runtime->run(g)`（执行 kernel）都能正常工作。 + +**验证**：`test_nativecpu_*` 三个测例通过即说明 dataMalloc 正确； +`test_allocator` 通过即说明 alloc/free 正确。 + +--- + +## 作业二：transpose 算子形状推导（⭐） + +**对应测例**：`test_transpose`、`test_nativecpu_transpose` +**涉及文件**：`src/operators/transpose.cc` + +### 2.1 原理 + +ONNX Transpose 算子有一个属性 `permute`：**输出的第 i 维来自输入的第 `permute[i]` 维**。 +例如输入形状 `[1,2,3,4]`，`permute = {0,2,1,3}`，则输出形状为 `[1,3,2,4]`。 +`transposePermute` 成员在构造函数中已初始化（空 permute 时默认为恒等置换）。 + +**ONNX 标准参考**： + +### 2.2 实现 + +```cpp +optional> TransposeObj::inferShape(const TensorVec &inputs) +{ + const auto A = inputs[0]; + auto input_dim = A->getDims(); + auto output_dim = input_dim; + int rank = A->getRank(); + + // =================================== 作业 =================================== + // permute 的含义：输出的第 i 维来自输入的第 permute[i] 维， + // 即 output_dim[i] = input_dim[permute[i]]。 + for (int i = 0; i < rank; ++i) + output_dim[i] = input_dim[transposePermute[i]]; + return {{output_dim}}; + // =================================== 作业 =================================== +} +``` + +**验证**：`test_transpose` 的三个用例（含恒等置换）都期望 `getOutput()->getDims()` +等于置换后的形状；`test_nativecpu_transpose` 进一步验证数据层面也正确。 + +--- + +## 作业三：clip 算子形状推导（⭐） + +**对应测例**：`test_clip` +**涉及文件**：`src/operators/unary.cc` + +### 3.1 原理 + +Clip 把每个元素限制在 `[min, max]` 区间内，**不改变形状**，输出形状与输入完全一致。 + +**ONNX 标准参考**： + +### 3.2 实现 + +```cpp +optional> ClipObj::inferShape(const TensorVec &inputs) +{ + // =================================== 作业 =================================== + // clip 是按元素把数值限制在 [min, max] 区间内，不改变 tensor 的形状， + // 输出形状与输入形状完全一致。 + return {{inputs[0]->getDims()}}; + // =================================== 作业 =================================== +} +``` + +**注意返回类型**：`optional>`，最外层 `{}` 是 optional， +`{{...}}` 是 vector（含一个 Shape）。不要写错括号层级。 + +--- + +## 作业四：cast 算子形状推导 & 数据类型推导（⭐⭐） + +**对应测例**：`test_cast` +**涉及文件**：`src/operators/unary.cc` + +### 4.1 原理 + +Cast 把元素从一种数据类型转为另一种，**形状不变**，**数据类型改变**。 +`CastType` 到 `DataType` 的映射已经由现成的 `getOutputDataType()` 实现 +（`src/operators/unary.cc` 底部的 switch 语句），我们只需把它接进来。 + +**ONNX 标准参考**： + +### 4.2 实现 + +```cpp +vector CastObj::inferDataType(const TensorVec &inputs) const +{ + // =================================== 作业 =================================== + // cast 只有一个输出，数据类型由 castType 决定。 + // getOutputDataType() 已经把 CastType 到 DataType 的映射实现好了。 + return {getOutputDataType()}; + // =================================== 作业 =================================== +} + +optional> CastObj::inferShape(const TensorVec &inputs) +{ + // =================================== 作业 =================================== + // cast 只改变数据类型，不改变形状，输出形状与输入一致。 + return {{inputs[0]->getDims()}}; + // =================================== 作业 =================================== +} +``` + +**解释：** + +- `inferDataType` 返回 `vector`，**第 i 个元素是第 i 个输出的数据类型**。 + cast 只有一个输出，所以返回 `{getOutputDataType()}`。默认实现 + （`src/core/operator.cc:74`）是"输出类型 = 输入类型"，对 cast 不适用，必须重写。 +- `checkValid` 中会调用 `inferDataType` 来创建输出 tensor，所以测例里的 + `op->getOutDType()` 才等于 `Float16`。 + +--- + +## 作业五：concat 算子形状推导（⭐⭐） + +**对应测例**：`test_concat`、`test_nativecpu_concat` +**涉及文件**：`src/operators/concat.cc` + +### 5.1 原理 + +Concat 把多个输入在 `dim` 维度上拼接： +- 其它维度保持不变； +- `dim` 维的大小 = 所有输入在该维大小之和。 + +例如 `[1,3,2,4]` 与 `[1,3,2,5]` 在 dim=3 上拼接，得到 `[1,3,2,9]`。 +构造函数里已经用 `get_real_axis` 把负轴归一化到了 `[0, rank)`。 + +**ONNX 标准参考**： + +### 5.2 实现 + +```cpp +optional> ConcatObj::inferShape(const TensorVec &inputs) { + Shape dims = inputs[0]->getDims(); + auto rank = inputs[0]->getRank(); + + // =================================== 作业 =================================== + IT_ASSERT(rank >= 1); + size_t concatDim = (size_t)get_real_axis(dim, rank); // 归一化 dim（防御负轴） + dims[concatDim] = 0; // 先清零拼接维 + for (auto &input : inputs) { + IT_ASSERT(input->getRank() == rank); // 所有输入维度数必须一致 + for (size_t i = 0; i < rank; ++i) + if (i != concatDim) // 除拼接维外，其余维度必须相同 + IT_ASSERT(input->getDims()[i] == dims[i], + "Concat: mismatch on dim " + std::to_string(i)); + dims[concatDim] += input->getDims()[concatDim]; // 累加各输入在该维的大小 + } + // =================================== 作业 =================================== + + return {{dims}}; +} +``` + +**验证**：`test_concat` 断言输出为 `{1,3,2,9}`；`test_nativecpu_concat` 验证 +三个输入拼接后的数据与期望完全一致。 + +**健壮性说明**：ONNX 要求所有输入**除 dim 维外形状完全相同**。上面的实现除了 +检查 `rank` 一致外，还逐一断言了每个非拼接维与其他输入相同（`i != concatDim`）。 +如果只检查 rank 而不检查其它维度，非法的输入（如 `[1,3,2,4]` 与 `[1,3,2,5]` +在 dim=1 上拼接）会被静默接受并产生错误结果，因此这一步检查建议保留。 + +**类型陷阱**：`getRank()` 返回 `size_t`（`TensorObj::getRank()` 返回 +`shape.size()`），而项目编译选项带有 `-Wall -Werror`（警告视为错误）， +所以循环变量必须用 `size_t`（如 `for (size_t i = 0; ...)`）并把 +`concatDim` 也转成 `size_t` 再比较；若写 `for (int i = 0; i < rank; ++i)`， +会触发 `-Wsign-compare` 直接编译失败。 + +--- + +## 作业六：双向广播（⭐⭐⭐） + +**对应测例**：`test_element_wise`、`test_nativecpu_elementwise`、`test_matmul` + +**涉及文件**：`src/utils/operator_utils.cc`（`infer_broadcast`） + +### 6.1 原理 + +NumPy/ONNX 的广播规则（从后往前对齐）： + +1. 两个 shape 从**最后一个维度**开始逐维比较； +2. 维度数量不同时，较短的 shape 在前面补 1； +3. 每个维度上，若两个大小相等、或其中一个是 1，则可以广播，结果取**较大的那个**； +4. 若两个都大于 1 且不相等，则不能广播（报错）。 + +例如 `{1,2,2,3,1}` 与 `{2,1,1}` 广播为 `{1,2,2,3,2}`（对齐后： +`{1,2,2,3,1}` vs `{1,1,2,1,1}` → `{1,2,2,3,2}`）。 + +**ONNX 标准参考**： + +### 6.2 实现 + +```cpp +Shape infer_broadcast(const Shape &A, const Shape &B) { + + // =================================== 作业 =================================== + Shape ans; + auto itA = A.rbegin(), itB = B.rbegin(); // 反向迭代器：从后往前比较 + while (itA != A.rend() || itB != B.rend()) + { + int dimA = (itA != A.rend()) ? *itA++ : 1; // 短的 shape 视为补 1 + int dimB = (itB != B.rend()) ? *itB++ : 1; + IT_ASSERT(dimA == dimB || dimA == 1 || dimB == 1, + "Broadcast failed: " + vecToString(A) + " vs " + vecToString(B)); + ans.emplace_back(std::max(dimA, dimB)); // 每维取较大者 + } + // 上面是从后往前生成的，翻转回来即得到从左到右的形状 + std::reverse(ans.begin(), ans.end()); + return ans; + // =================================== 作业 =================================== +} +``` + +**逐段解释：** + +- 用**反向迭代器**天然实现"从后往前对齐"；当一个 shape 较短时，对应的迭代器先到 + `rend()`，此时该维按 1 处理（等价于前面补 1）。 +- `std::max(dimA, dimB)` 实现"取较大者"：相等取本身，有一方为 1 取另一方。 +- 因为是从后往前 push 的，最后要 `std::reverse` 一下恢复从左到右的顺序。 +- 需要 `#include `（文件顶部已加）。 + +**验证**：`test_element_wise` 的 5 个用例覆盖了各种广播组合（标量、单维、多 1、 +对齐补位等）；`test_nativecpu_elementwise` 验证广播后的实际计算数据正确； +`test_matmul` 的批量维度广播也依赖本函数。 + +--- + +## 作业七：矩阵乘形状推导（⭐⭐⭐） + +**对应测例**：`test_matmul` +**涉及文件**：`src/operators/matmul.cc` + +### 7.1 原理 + +ONNX MatMul 语义（行主序）： + +- A 的形状记为 `[batch..., m, k']`，B 记为 `[batch..., k'', n]`； +- `transA = true` 表示 A 参与乘法前**交换最后两个维度**，即转置前的 A 形状是 + `[batch..., k, m]`； +- `transB = true` 表示 B 参与乘法前交换最后两个维度，即转置前的 B 形状是 + `[batch..., n, k]`； +- 两个矩阵的 k 维必须相等； +- 输出形状 = 批量维（A、B 的批量维双向广播）+ `[m, n]`。 + +**ONNX 标准参考**： + +**⚠️ 最容易写反的地方：m/k 的取值** + +| transA | A 转置前形状 | m | k | +| --- | --- | --- | --- | +| false | `[..., m, k]` | `dimsA[rankA-2]` | `dimsA[rankA-1]` | +| true | `[..., k, m]` | `dimsA[rankA-1]` | `dimsA[rankA-2]` | + +| transB | B 转置前形状 | n | k | +| --- | --- | --- | --- | +| false | `[..., k, n]` | `dimsB[rankB-1]` | `dimsB[rankB-2]` | +| true | `[..., n, k]` | `dimsB[rankB-2]` | `dimsB[rankB-1]` | + +直观记忆：**不转置时，m/k 分别来自倒数第二维/最后一维；转置时正好对调**。 + +### 7.2 实现 + +```cpp +optional> MatmulObj::inferShape(const TensorVec &inputs) +{ + // =================================== 作业 =================================== + const auto A = inputs[0], B = inputs[1]; + auto dimsA = A->getDims(); + auto dimsB = B->getDims(); + auto rankA = A->getRank(); + auto rankB = B->getRank(); + IT_ASSERT(rankA >= 2 && rankB >= 2); + + // 解析出 m、n、k（见上表，trans 时 m/k 的来源对调） + m = transA ? dimsA[rankA - 1] : dimsA[rankA - 2]; + k = transA ? dimsA[rankA - 2] : dimsA[rankA - 1]; + n = transB ? dimsB[rankB - 2] : dimsB[rankB - 1]; + IT_ASSERT(k == (transB ? dimsB[rankB - 1] : dimsB[rankB - 2]), + "Matmul: inner dimensions mismatch"); + + // 批量维度做双向广播（复用作业六的 infer_broadcast） + Shape batchA(dimsA.begin(), dimsA.end() - 2); + Shape batchB(dimsB.begin(), dimsB.end() - 2); + auto batch = infer_broadcast(batchA, batchB); + + // 拼接得到输出形状 + Shape dimsC = batch; + dimsC.emplace_back(m); + dimsC.emplace_back(n); + return {{dimsC}}; + // =================================== 作业 =================================== +} +``` + +**解释：** + +- `m`、`n`、`k` 是 `MatmulObj` 的成员变量（头文件里声明了），这里顺便赋值， + 供 `toString()` 打印使用。 +- 批量维是 `dimsA` 去掉最后两维的部分；两个批量维形状不同时按广播规则对齐。 + 测例 `{1,2,3,5} × {1,1,5,2}` → 批量 `{1,2} 广播 {1,1}` → `{1,2}`， + 输出 `{1,2,3,2}`，与期望一致。 +- 本作业依赖作业六（`infer_broadcast`），需要 + `#include "utils/operator_utils.h"`。 + +--- + +## 作业八：简单图优化规则实现（⭐⭐⭐⭐） + +**对应测例**：`test_graph` +**涉及文件**：`src/core/graph.cc`（`optimize`）、`src/core/operator.cc`（被调用的接口） + +### 8.1 题目要求 + +1. **去除冗余算子**：两个相邻 transpose 互为逆操作（permute 互逆）时，两者可以 + 一起删除，中间结果直接跳过； +2. **合并算子**：matmul 的输入若来自"仅交换最后两维"的 transpose，可以把它 + 融合进 matmul 的 `transA`/`transB` 属性。 + +### 8.2 先看测例想要什么（test_graph 的图） + +```text +i1 --T{0,1,3,2}--> t1 --T{0,1,3,2}--> t2 ─┐ + ├──> Matmul ──> o +i2 --T{0,1,3,2}--> t3 ─────────────────────┘ +``` + +测例的断言： + +- 优化后只有 **1 个算子**（MatMul，`OpType::MatMul == 7`）； +- 只剩 **3 个 tensor**（i1、i2、o）； +- matmul 输入为 i1、i2（guid 分别为 2、3，因为 graph 占 guid 1）； +- `transA == false`、`transB == true`。 + +即期望结果：`Matmul(i1, i2, o, transA=false, transB=true)`。 + +**注意顺序**：必须先执行规则 1（消除 `t1/t2` 之间的互逆对），再执行规则 2。 +如果反过来，规则 2 会先把 `t2`（来自 t1 的 transpose）融合成 `transA=true`， +之后规则 1 虽然删掉了互逆对，但 `transA` 仍残留为 true，导致断言失败。 + +### 8.3 完整实现 + +```cpp +void GraphObj::optimize() +{ + // =================================== 作业 =================================== + // 辅助函数 1：判断两个 permute 是否互为逆置换（perm1 ∘ perm2 = identity） + auto isInverse = [](const std::vector &p1, const std::vector &p2) { + if (p1.size() != p2.size()) + return false; + for (size_t i = 0; i < p1.size(); ++i) + if (p2[(size_t)p1[i]] != (int)i) + return false; + return true; + }; + // 辅助函数 2：判断 permute 是否只交换最后两个维度（如 {0,1,3,2}） + auto isSwapLastTwo = [](const std::vector &perm) { + size_t r = perm.size(); + if (r < 2) + return false; + for (size_t i = 0; i < r; ++i) + { + int expect = (int)i; // 默认保持原位 + if (i == r - 2) // 倒数第二维映射到最后一位 + expect = (int)(r - 1); + else if (i == r - 1) // 最后一位映射到倒数第二维 + expect = (int)(r - 2); + if (perm[i] != expect) + return false; + } + return true; + }; + + // ==================== 规则 1：去除冗余的互逆 transpose 对 ==================== + bool changed = true; + while (changed) + { + changed = false; + for (auto &op : ops) + { + if (op->getOpType() != OpType::Transpose) + continue; + auto trans = as(op); // mid -> out + auto mid = trans->getInputs(0); + auto pred = mid->getSource(); // input -> mid + if (!pred || pred->getOpType() != OpType::Transpose) + continue; + // 安全条件：mid 只被 trans 这一个算子消费，否则不能删除 pred + if (mid->getTargets().size() != 1) + continue; + if (!isInverse(as(pred)->getPermute(), + trans->getPermute())) + continue; + + auto input = pred->getInputs(0); + auto out = trans->getOutput(0); + // 先记住 out 是否有消费者（决定它是否还能删除） + bool outHasConsumers = !out->getTargets().empty(); + // 把 out 的所有消费者（后继算子）的输入从 out 换成 input + for (auto &succ : out->getTargets()) + { + succ->replaceInput(out, input); + out->removeTarget(succ); + input->addTarget(succ); + succ->removePredecessors(trans); + trans->removeSuccessors(succ); + } + input->removeTarget(pred); + mid->removeTarget(trans); + trans->removePredecessors(pred); + pred->removeSuccessors(trans); + mid->setSource(nullptr); + out->setSource(input->getSource()); + // 从图中移除冗余的算子和中间 tensor + removeOperator(trans); + removeOperator(pred); + removeTensor(mid); + // out 原本有消费者且已全部改接到 input，可以一并删除； + // 若 out 是图的输出（本来就没有消费者），则予以保留。 + if (outHasConsumers) + removeTensor(out); + changed = true; + break; // 重新扫描 + } + } + + // ============ 规则 2：把输入上的 transpose 融合进 matmul 的 trans 属性 ============ + for (auto &op : ops) + { + if (op->getOpType() != OpType::MatMul) + continue; + auto matmul = as(op); + for (int i = 0; i < 2; ++i) + { + auto in = matmul->getInputs(i); + auto src = in->getSource(); + if (!src || src->getOpType() != OpType::Transpose) + continue; + // 安全条件：该 transpose 只被 matmul 一个算子消费 + if (in->getTargets().size() != 1) + continue; + if (!isSwapLastTwo(as(src)->getPermute())) + continue; + auto trans = as(src); + auto newIn = trans->getInputs(0); + if (i == 0) + matmul->setTransA(true); + else + matmul->setTransB(true); + matmul->replaceInput(in, newIn); + in->removeTarget(matmul); + newIn->addTarget(matmul); + newIn->removeTarget(trans); + matmul->removePredecessors(trans); + trans->removeSuccessors(matmul); + in->setSource(nullptr); + removeOperator(trans); + removeTensor(in); + } + } + // =================================== 作业 =================================== +} +``` + +### 8.4 逐段解释 + +**辅助函数** + +- `isInverse(p1, p2)`：判断 `p1` 和 `p2` 是否互逆。数学上两个置换互逆当且仅当 + 它们的复合是恒等置换，即对每个维度 i，`p2[p1[i]] == i`。 + `{0,1,3,2}` 是自逆的（交换两个维度两次等于没换），所以测例中的 + `T{0,1,3,2} → T{0,1,3,2}` 是一对互逆 transpose。 +- `isSwapLastTwo(perm)`：判断 permute 是否只交换最后两个维度——前 `r-2` 个维度 + 保持原位，第 `r-2` 维映射到 `r-1`，第 `r-1` 维映射到 `r-2`。 + +**规则 1 的删除流程** + +以 `pred(input -> mid)`、`trans(mid -> out)` 互逆为例。**删除前的图结构**如下 +（实线为数据流，括号内是每个 tensor 的 `source` / `targets`）： + +```text +[input] ──(pred)──► [mid] ──(trans)──► [out] ──────────► [succ1] [succ2] + source:null source:pred source:trans （succ 是 out 的消费者） + targets:[pred] targets:[trans] targets:[succ1,succ2] +``` + +op 侧还有 `pred.succ=[trans]`、`trans.pred=[pred]`、`succ1.pred=[trans]` 等 +前驱/后继关系，与 tensor 侧一一对应。整个删除分三步： + +**第 1 步：数据流改道**（循环遍历 `out` 的每个消费者 `succ`） + +```text + ┌──────────── 旧路径（即将废弃）────────────┐ + ▼ │ +[input] ──(pred)──► [mid] ──(trans)──► [out] ──► [succ1] [succ2] + └────────────── 新路径（replaceInput）─────────► ▲ + │ +``` + +- `succ->replaceInput(out, input)`：只改 succ 的输入列表（数据流改道）； +- `out->removeTarget(succ)` / `input->addTarget(succ)`：同步 tensor 侧 `targets`； +- `succ->removePredecessors(trans)` / `trans->removeSuccessors(succ)`：同步 op 侧关系。 + +**第 2 步：清理与被删算子相关的残留关系** + +| 操作 | 作用 | +| --- | --- | +| `input->removeTarget(pred)` | input 的消费者列表移除 pred | +| `mid->removeTarget(trans)` | mid 的消费者列表移除 trans | +| `trans->removePredecessors(pred)` | trans 的前驱移除 pred | +| `pred->removeSuccessors(trans)` | pred 的后继移除 trans | +| `mid->setSource(nullptr)` | mid 的生产者（pred）已删，置空 | +| `out->setSource(input->getSource())` | out 的生产者改为 input 的生产者 | + +**第 3 步：从图中移除冗余对象**（`removeOperator(trans/pred)`、`removeTensor(mid)`）， +**删除后的图结构**： + +```text +[input] ──────────────────────────────────► [succ1] [succ2] + source:null（图输入不变） targets:[succ1,succ2] +``` + +几点说明： + +1. 中间 tensor `mid`（t1）一定可以删；`out`（t2）只有当它**原来有消费者** + （已被重连到 input）时才删——如果 `out` 是图的输出（本来就没有消费者）， + 必须保留，否则图的输出会丢失。 +2. 删除前必须用 `mid->getTargets().size() == 1` 确认 `mid` 只被 `trans` 消费， + 否则 `pred` 不能删。 +3. `source/targets`、`predecessors/successors` 存的是 `weak_ptr`，**任何一侧 + 残留指向被删算子的引用都会变成失效的 weak_ptr**，打印图时直接段错误或抛 + `bad_weak_ptr`——这正是第 1、2 步必须"同步维护双向关系"的原因。 + +`while (changed)` 循环保证链式冗余（如三个以上的 transpose 链）也能被消除： +每轮删除一对后重新扫描，直到找不到为止。 + +**规则 2 的融合流程** + +matmul 的两个输入分别处理：若输入 `in` 的 `source` 是 transpose 且 +`isSwapLastTwo(perm)` 成立，则： + +**融合的本质是"语义等价改写"**：matmul 的 `transA`/`transB` 属性语义就是 +"参与乘法前交换最后两个维度"，与"只交换最后两维的 transpose"完全等价： + +```text +Matmul( T(A), B ) ≡ Matmul( A, B, transA=true ) +Matmul( A, T(B) ) ≡ Matmul( A, B, transB=true ) +``` + +**融合前的图结构**（以测例图的 B 输入为例，实线为数据流，括号内为 +`source` / `targets`）： + +```text +[i2] ──(T10: perm={0,1,3,2})──► [t3] ──────► Matmul(A=i1, B=t3) ──► [o] + source:null source:T10 source:Matmul + targets:[T10] targets:[Matmul] targets:[] +``` + +**三个前置条件（缺一不可）**： + +1. `in->getSource()` 存在且是 Transpose——输入确实由 transpose 产生； +2. `in->getTargets().size() == 1`——该 transpose 只被本 matmul 消费， + 否则删除 transpose 会让其他消费者输入悬空； +3. `isSwapLastTwo(perm)`——`transA`/`transB` 只能表达"交换最后两维" + 这一种转置，其他 permute 模式无法等价表达。 + +**融合的四个动作**（设置属性 → 数据流改道 → 同步关系 → 删除）： + +```cpp +matmul->setTransB(true); // ① 属性置位（处理 A 时用 setTransA） +matmul->replaceInput(t3, i2); // ② 数据流改道 +removeOperator(trans); // ③ 删除 transpose +removeTensor(t3); // ④ 删除中间 tensor +``` + +```text + ┌──────────────── 旧路径（T10 已删除）────────────┐ + ▼ │ +[i2] ──(T10)──► [t3] ──────► Matmul(A=i1, B=t3) ──► [o] + └──────────── 新路径（replaceInput）────────────► ▲ +``` + +① 属性置位后，matmul 计算时自行交换 `i2` 的最后两维——原来由 T10 做的事 +改由 matmul 内部完成，`m/n/k` 与输出形状均不变。 + +**同步维护双向关系**（与规则 1 相同的清单，防止残留失效的 `weak_ptr`）： + +| 操作 | 作用 | +| --- | --- | +| `in->removeTarget(matmul)` | t3 的消费者列表移除 matmul | +| `newIn->addTarget(matmul)` | i2 的消费者列表新增 matmul | +| `newIn->removeTarget(trans)` | i2 的消费者列表移除 T10（**易漏！**） | +| `matmul->removePredecessors(trans)` | matmul 的前驱移除 T10 | +| `trans->removeSuccessors(matmul)` | T10 的后继移除 matmul | +| `in->setSource(nullptr)` | t3 的生产者清空（随后删除） | + +**融合后的图结构**： + +```text +[i2] ──────────────────────────────► Matmul(A=i1, B=i2, transB=true) ──► [o] + source:null（图输入不变） source:Matmul + targets:[Matmul] targets:[] +``` + +**为什么必须放在规则 1 之后**：若先做规则 2，matmul 的 A 输入 `t2` 的 source +是 `T9`（互逆对的一员），会被误融合成 `transA=true`；随后规则 1 删掉互逆对时 +`transA` 残留为 true，而期望值是 `false`。先做规则 1 后，matmul 的 A 直接变成 +图输入 `i1`（无 source 可融合），B 只剩独立的 T10 可融合， +`transA=false, transB=true` 才正确。 + +**为什么规则 2 不需要 while 循环**：融合发生在 matmul 的输入上，每融合一个就 +删除一个 transpose，不会产生新的可融合模式，单遍遍历即可。 + +### 8.5 测例逐步演算 + +初始图（guid 顺序：graph=1，i1=2，i2=3，t1=4，t2=5，t3=6，o=7）： + +```text +T8: i1 -> t1, perm={0,1,3,2} +T9: t1 -> t2, perm={0,1,3,2} +T10: i2 -> t3, perm={0,1,3,2} +M11: Matmul(t2, t3) -> o +``` + +- **规则 1**：找到 `T9`（trans），`t1` 的 source 是 `T8`（pred），两者 permute + 相同故互逆，`t1` 只被 T9 消费 → 删除 T8、T9。`t2` 的唯一消费者 M11 的输入 + 从 `t2` 改为 `i1`。删除中间 tensor `t1`、`t2`。图变为： + ```text + T10: i2 -> t3 + M11: Matmul(i1, t3) -> o + ``` + 第二轮扫描：`t3` 的 source 是 T10，而 `i2` 没有 source（图输入）， + `pred` 不存在 → 无冗余对，退出循环。 +- **规则 2**：M11 的输入 A=`i1`（无 source，跳过）；输入 B=`t3`，其 source 是 + T10 且 perm `{0,1,3,2}` 是 swap-last-two → `setTransB(true)`，输入换成 `i2`， + 删除 T10 和 `t3`。 +- **结果**：`Matmul(i1, i2, o, false, true)`，1 个算子、3 个 tensor，全部断言通过。 + +--- + +## 9. 自测与常见问题 + +### 9.1 如何自测 + +```bash +make build # 编译 +make test-cpp # 跑全部 11 个测例 +``` + +也可以单独跑某个测例（定位问题更快）： + +```bash +cd build/Release && ./test_allocator # 或 ./test_graph / ./test_matmul ... +``` + +用 Debug 模式（`make TYPE=Debug build`）编译可以获得更清晰的调试信息； +用 gdb 加 `catch throw` 可以捕获测例里抛出的异常。 + +### 9.2 常见问题排查清单 + +| 现象 | 原因 | 对策 | +| --- | --- | --- | +| 多个测例抛空描述异常（`bad_alloc`） | matmul 里 m/n/k 推导写反，`Shape` 构造出巨大量 | 对照 7.1 的表格逐项检查 trans 与 m/k 的对应关系 | +| `testAllocWithEndFreeBlock` 的 offsetD != offsetC | alloc 缺少"末尾空闲块扩展"逻辑 | 补上 1.3 中的末尾扩展分支 | +| 二次 `alloc` 断言 `this->ptr == nullptr` 失败 | dataMalloc 边 alloc 边调 `getPtr()` | 改为两阶段：先规划偏移量，最后统一 `getPtr()` 绑 Blob | +| nativecpu 测例数据算错（尤其是 concat） | 算子的输出复用了本算子输入的内存 | 调整顺序：先分配输出，再回收输入 | +| `test_graph` 段错误 / `bad_weak_ptr` | 删除算子/重连时没有同步维护 `targets`、`predecessors/successors` | 按 8.4 的清单补齐双向关系维护 | +| `test_graph` 的 `transA` 断言失败 | 规则 1、规则 2 执行顺序颠倒 | 必须先消除互逆对，再融合 matmul | +| cast 测例的 `getOutDType()` 不对 | `inferDataType` 没重写（返回了输入类型） | 返回 `{getOutputDataType()}` | +| 编译报 `std::reverse` 等未定义 | 缺头文件 | `operator_utils.cc` 需要 `#include ` | + +### 9.3 进阶思考 + +- **分配策略**：first-fit 之外还有 best-fit（找最接近 size 的块）、worst-fit 等， + 各有优劣；真实框架（如 InfiniTensor 的 CUDA 后端）还会考虑对齐、跨设备等。 +- **内存复用**：`dataMalloc` 里的引用计数复用是简化版的"生命周期分析"， + 真实编译器会对整个图做全局的活跃区间（liveness）分析再统一着色。 +- **图优化**：规则 2 只是"算子融合"的一个例子，工业级实现还有常量折叠、 + 算子重排、kernel 融合（fuse）、死代码消除等。 diff --git "a/docs/\351\200\232\350\277\207\346\210\252\345\233\276.png" "b/docs/\351\200\232\350\277\207\346\210\252\345\233\276.png" new file mode 100644 index 00000000..f9e5fea2 Binary files /dev/null and "b/docs/\351\200\232\350\277\207\346\210\252\345\233\276.png" differ diff --git "a/docs/\351\241\271\347\233\256\351\203\250\347\275\262.md" "b/docs/\351\241\271\347\233\256\351\203\250\347\275\262.md" index 5690349b..2cb9c9a9 100644 --- "a/docs/\351\241\271\347\233\256\351\203\250\347\275\262.md" +++ "b/docs/\351\241\271\347\233\256\351\203\250\347\275\262.md" @@ -27,9 +27,13 @@ sudo apt install make # mac 使用Homebrew安装 brew install make ``` - ### 构建命令 + 配置好上述环境后，进入项目目录后可以通过以下命令进行构建。 +> 注意：`3rd-party/googletest` 是 git submodule，`make build` 会自动执行 +> `git submodule update --init --recursive` 初始化；若网络不佳导致子模块 +> 拉取失败，可手动执行该命令后重试。 + - `make`/`make build`: 构建整个项目; - `make test-cpp`: 构建项目后执行测例; - `make clean`：清理生成文件 \ No newline at end of file diff --git a/include/core/allocator.h b/include/core/allocator.h index 002601d2..59a84ec2 100644 --- a/include/core/allocator.h +++ b/include/core/allocator.h @@ -27,7 +27,14 @@ namespace infini { // TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并 // HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小 // =================================== 作业 =================================== - + // 踩坑（详见 docs/作业完成说明.md 作业一）： + // - alloc/free 只是"模拟"分配（函数开头有 IT_ASSERT(ptr == nullptr)）， + // 真实内存必须在全部规划完成之后通过 getPtr() 一次性分配； + // - alloc 需要支持"末尾空闲块扩展"（地址最大的空闲块延伸到 peak 时， + // 可以从该块起始地址继续向外扩展），否则 test_allocator 的 + // testAllocWithEndFreeBlock 无法通过； + // - free 时记得与前一块、后一块做合并（先向前、再向后），避免碎片。 + std::map freeBlocks; public: Allocator(Runtime runtime); diff --git a/src/core/allocator.cc b/src/core/allocator.cc index ff593aef..18c75554 100644 --- a/src/core/allocator.cc +++ b/src/core/allocator.cc @@ -1,4 +1,7 @@ #include "core/allocator.h" +#include +#include +#include #include namespace infini @@ -32,8 +35,42 @@ namespace infini // =================================== 作业 =================================== // TODO: 设计一个算法来分配内存，返回起始地址偏移量 // =================================== 作业 =================================== + // - 推荐 first-fit：从低地址开始扫描 free list，找到第一个足够大的块； + // - 块比需要的大时，把剩余部分切回 free list； + // - 找不到合适块时，在 peak 处追加；若地址最大的空闲块恰好延伸到 + // peak（first + second == peak），从它的起始地址向外扩展（见头文件踩坑）； + // - 分配成功记得更新 used、peak（peak = max(peak, used)）。 + for(auto& it : freeBlocks) { + if(it.second >= size) { + size_t addr = it.first; + if(it.second > size) { + freeBlocks[addr + size] = it.second - size; + } + freeBlocks.erase(it.first); + used += size; + peak = std::max(peak, used); + return addr; + } + } + // 如果地址最大的空闲块紧贴着内存末尾 + // 为了复用内存，就从该内存块开始向后拓展 + if(freeBlocks.empty() != true) { + auto last = std::prev(freeBlocks.end()); + if(last->first + last->second == peak) { + last->second += size; + size_t addr = last->first; + used += size; + peak = addr + size; + freeBlocks.erase(last); + return addr; + } + } + + size_t addr = peak; + used += size; + peak += size; - return 0; + return addr; } void Allocator::free(size_t addr, size_t size) @@ -44,6 +81,25 @@ namespace infini // =================================== 作业 =================================== // TODO: 设计一个算法来回收内存 // =================================== 作业 =================================== + // - 把 [addr, addr + size) 放回 free list（used 相应减少）； + // - 先与前一块合并（prev->first + prev->second == addr）， + // 再与后一块合并（it->first + it->second == next->first），顺序不能反。 + auto it = freeBlocks.emplace(addr, size).first; + used -= size; + + if(it != freeBlocks.begin()) { + auto prev = std::prev(it); + if(prev->first + prev->second == addr) { + prev->second += size; + freeBlocks.erase(it); + it = prev; + } + } + auto next = std::next(it); + if(next != freeBlocks.end() && it->first + it->second == next->first) { + it->second += next->second; + freeBlocks.erase(next); + } } void *Allocator::getPtr() diff --git a/src/core/graph.cc b/src/core/graph.cc index 3a906370..54a1af1d 100644 --- a/src/core/graph.cc +++ b/src/core/graph.cc @@ -1,7 +1,17 @@ #include "core/graph.h" +#include "core/blob.h" +#include "core/ref.h" +#include "core/runtime.h" +#include "operators/matmul.h" +#include "operators/transpose.h" #include +#include +#include +#include #include #include +#include +#include namespace infini { @@ -106,6 +116,136 @@ namespace infini // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除） // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去） // =================================== 作业 =================================== + // 踩坑（详见 docs/作业完成说明.md 作业八）： + // - 执行顺序：必须先做规则 1（消除互逆 transpose 对），再做规则 2 + // （融合进 matmul），否则本应被抵消的 transpose 会被错误融合成 transA； + // - 删除算子/重连输入时，必须同步维护双向关系（tensor 的 source/targets + // 与 op 的 predecessors/successors），否则会残留失效的 weak_ptr， + // 打印图时崩溃（段错误 / bad_weak_ptr）； + // - 两个 permute 互逆 ⇔ 复合为恒等置换：p2[p1[i]] == i； + // - 只交换最后两维 ⇔ 前 rank-2 维不动，最后两维互换； + // - 删除中间 tensor 前检查它是否还被其他算子消费（targets 非空则保留）。 + auto isInverse = [](const std::vector &p1, const std::vector &p2) { + if(p1.size() != p2.size()) { + return false; + } + for(size_t i = 0; i < p1.size(); i++) { + if(p2[(size_t)p1[i]] != (int)i) { + return false; + } + } + return true; + }; + auto isSwapLastTwo = [](const std::vector &perm) { + size_t r = perm.size(); + if (r < 2) + return false; + for (size_t i = 0; i < r; ++i) + { + int expect = (int)i; // 默认保持原位 + if (i == r - 2) // 倒数第二维映射到最后一位 + expect = (int)(r - 1); + else if (i == r - 1) // 最后一位映射到倒数第二维 + expect = (int)(r - 2); + if (perm[i] != expect) + return false; + } + return true; + }; + // ==================== 规则 1：去除冗余的互逆 transpose 对 ==================== + bool changed = true; + while (changed){ + changed = false; + for (auto &op : ops){ + if (op->getOpType() != OpType::Transpose) { + continue; + } + auto trans = as(op); // mid -> out + auto mid = trans->getInputs(0); + auto pred = mid->getSource(); // input -> mid 获取mid的生产者 + if (!pred || pred->getOpType() != OpType::Transpose){ + continue; + } + // 安全条件：mid 只被 trans 这一个算子消费，否则不能删除 pred + if (mid->getTargets().size() != 1) { + continue; + } + // 如果两个permute不互为逆置换，continue + if (!isInverse(as(pred)->getPermute(), + trans->getPermute())) { + continue; + } + + auto input = pred->getInputs(0); + auto out = trans->getOutput(0); + // 先记住 out 是否有消费者（决定它是否还能删除） + bool outHasConsumers = !out->getTargets().empty(); + // 把 out 的所有消费者（后继算子）的输入从 out 换成 input + for (auto &succ : out->getTargets()) { + succ->replaceInput(out, input); + out->removeTarget(succ); + input->addTarget(succ); + succ->removePredecessors(trans); + trans->removeSuccessors(succ); + } + input->removeTarget(pred); + mid->removeTarget(trans); + trans->removePredecessors(pred); + pred->removeSuccessors(trans); + + removeOperator(trans); + removeOperator(pred); + removeTensor(mid); + + // out 原本有消费者且已全部改接到 input，可以一并删除； + // 若 out 是图的输出（本来就没有消费者），则予以保留。 + if (outHasConsumers) { + removeTensor(out); + } + changed = true; + break; // 重新扫描 + } + } + // ============ 规则 2：把输入上的 transpose 融合进 matmul 的 trans 属性 ============ + for (auto &op : ops) + { + if (op->getOpType() != OpType::MatMul) { + continue; + } + auto matmul = as(op); + for (int i = 0; i < 2; ++i) + { + auto in = matmul->getInputs(i); + auto src = in->getSource(); + if (!src || src->getOpType() != OpType::Transpose) { + continue; + } + // 安全条件：该 transpose 只被 matmul 一个算子消费 + if (in->getTargets().size() != 1) { + continue; + } + if (!isSwapLastTwo(as(src)->getPermute())) { + continue; + } + auto trans = as(src); + auto newIn = trans->getInputs(0); + if (i == 0) { + matmul->setTransA(true); + } + else{ + matmul->setTransB(true); + } + matmul->replaceInput(in, newIn); + in->removeTarget(matmul); + newIn->addTarget(matmul); + newIn->removeTarget(trans); + matmul->removePredecessors(trans); + trans->removeSuccessors(matmul); + in->setSource(nullptr); + removeOperator(trans); + removeTensor(in); + } + } } Tensor GraphObj::getTensor(int fuid) const @@ -152,7 +292,57 @@ namespace infini // TODO：利用 allocator 给计算图分配内存 // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存 // =================================== 作业 =================================== + // 踩坑（详见 docs/作业完成说明.md 作业一 step3）： + // - 必须分两阶段：先对每个 tensor 做"模拟"分配（alloc/free，只算偏移量）， + // 最后统一调用 allocator.getPtr() 拿真实内存基址再绑定 Blob。 + // 不能边 alloc 边 getPtr()（alloc 开头有 IT_ASSERT(ptr == nullptr)）； + // - 建议用引用计数做内存复用：记录每个 tensor 剩余消费者个数（targets 数）， + // 某个 op 规划完后其输入计数减一，归零即 free，供后面算子复用； + // - 每个 tensor 只分配一次（用一个集合记录已分配的 tensor，避免重复分配）； + // - 重要：算子的"输出"必须先于"输入回收"来分配， + // 否则输出会复用本算子输入的内存，而 kernel 执行时边读输入边写输出， + // 两者重叠会导致结果错误； + // - 绑定方式：make_ref(runtime, (char *)base + offset)。 + std::map refCount; // 每个 Tensor 的消费者总数 + std::map offsetMap; // 每个 Tensor 的偏移 + std::unordered_set allocated; // 已分配的 Tensor + for(auto &t: tensors) { + refCount[t] = t->getTargets().size(); + } + // 阶段一：模拟分配，规划每个 tensor 的偏移量 + for(auto& op : ops) { + // 为算子的输入分配内存（每个 tensor 只规划一次） + for(auto& t : op->getInputs()) { + if(allocated.find(t) == allocated.end()) { + offsetMap[t] = allocator.alloc(t->getBytes()); + allocated.insert(t); + } + } + + // 为算子的输出分配内存（先于输入回收，避免与输入重叠） + for(auto& t : op->getOutputs()) { + if(allocated.find(t) == allocated.end()) { + offsetMap[t] = allocator.alloc(t->getBytes()); + allocated.insert(t); + } + } + // 本算子的输入在本算子执行完后不再被它使用： + // 引用计数减一，归零的输入（最后一个消费者已完成）立即回收， + // 供后面的算子复用，从而降低峰值内存。 + for(auto& t : op->getInputs()) { + --refCount[t]; + if(refCount[t] == 0) { + allocator.free(offsetMap[t], t->getBytes()); + } + } + } + // 阶段二：真正分配内存，并按偏移量给每个 tensor 绑定 Blob + auto base = allocator.getPtr(); + for(auto &[t, offset] : offsetMap) { + t->setDataBlob(make_ref(runtime, (char*)base + offset)); + } + allocator.info(); } diff --git a/src/operators/concat.cc b/src/operators/concat.cc index d1963308..58023f16 100644 --- a/src/operators/concat.cc +++ b/src/operators/concat.cc @@ -1,5 +1,6 @@ #include "operators/concat.h" #include "utils/operator_utils.h" +#include namespace infini { ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim) @@ -17,6 +18,25 @@ optional> ConcatObj::inferShape(const TensorVec &inputs) { // TODO：修改 dims，返回正确的 concat 后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13 // =================================== 作业 =================================== + // 除 dim 维外其它维度保持不变，dim 维大小为所有输入在该维大小之和； + // 构造函数已把 dim 归一化到 [0, rank)，实现里再用 get_real_axis(dim, rank) + // 防御一次负轴更稳妥； + // 注意先 dims[dim] = 0 再累加，且所有输入的 rank 必须一致。 + // 踩坑：除 dim 维外，其它维度也必须逐一断言相等（ONNX 要求所有输入除 + // concat 维外形状相同），只检查 rank 是不够的。 + // 注意:上面 auto rank 推导出来的类型是 size_t 即unsigned long + size_t concatDim = get_real_axis(dim, rank); + dims[dim] = 0; + for(auto& input : inputs) { + IT_ASSERT(input->getRank() == rank); + for (size_t i = 0; i < rank; ++i) { + if (i != concatDim) { // 除拼接维外，其余维度必须相同 + IT_ASSERT(input->getDims()[i] == dims[i], + "Concat: mismatch on dim " + std::to_string(i)); + } + } + dims[concatDim] += (input->getDims())[concatDim]; + } return {{dims}}; } diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc index 7a16ca27..875d0373 100644 --- a/src/operators/matmul.cc +++ b/src/operators/matmul.cc @@ -1,4 +1,7 @@ #include "operators/matmul.h" +#include "core/common.h" +#include "core/tensor.h" +#include "utils/operator_utils.h" namespace infini { @@ -27,7 +30,39 @@ namespace infini // TODO：返回经过 matmul 操作后的 shape // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm // =================================== 作业 =================================== - return std::nullopt; + // 踩坑（详见 docs/作业完成说明.md 作业七）： + // - m/n/k 的取值最容易写反！记 A 转置前为 [..., m, k]， + // B 转置前为 [..., k, n]： + // transA=false：m=dimsA[rankA-2]，k=dimsA[rankA-1] + // transA=true ：m=dimsA[rankA-1]，k=dimsA[rankA-2] + // transB=false：n=dimsB[rankB-1]，k=dimsB[rankB-2] + // transB=true ：n=dimsB[rankB-2]，k=dimsB[rankB-1] + // （trans 时 m/k 的来源正好对调） + // - 校验两边的 k 相等； + // - 输出 = 批量维（A/B 去掉最后两维后做双向广播，复用 + // utils/operator_utils.h 的 infer_broadcast）拼接 [m, n]； + // - m/n/k 是成员变量，这里顺便赋值，供 toString() 使用。 + const auto A = inputs[0]; + const auto B = inputs[1]; + auto dimsA = A->getDims(); + auto dimsB = B->getDims(); + auto rankA = A->getRank(); + auto rankB = B->getRank(); + + m = transA ? dimsA[rankA - 1] : dimsA[rankA - 2]; + k = transA ? dimsA[rankA - 2] : dimsA[rankA - 1]; + n = transB ? dimsB[rankB - 2] : dimsB[rankB - 1]; + IT_ASSERT(k == (transB ? dimsB[rankB - 1] : dimsB[rankB - 2]), + "Matmul: inner dimensions mismatch"); + + Shape batchA(dimsA.begin(), dimsA.end() - 2); + Shape batchB(dimsB.begin(), dimsB.end() - 2); + auto batch = infer_broadcast(batchA, batchB); + + auto dimsC = batch; + dimsC.emplace_back(m); + dimsC.emplace_back(n); + return {{dimsC}}; } } // namespace infini \ No newline at end of file diff --git a/src/operators/transpose.cc b/src/operators/transpose.cc index faab2b69..f7947fa9 100644 --- a/src/operators/transpose.cc +++ b/src/operators/transpose.cc @@ -33,8 +33,13 @@ namespace infini // TODO：修改 output_dim，返回正确的 transpose 后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-21 // =================================== 作业 =================================== - - return std::nullopt; + // permute 的含义是"输出的第 i 维来自输入的第 permute[i] 维"， + // 即 output_dim[i] = input_dim[transposePermute[i]]； + // 返回类型是 optional>，返回 {{output_dim}} 即可。 + for(int i = 0; i < rank; i++) { + output_dim[i] = input_dim[transposePermute[i]]; + } + return {{output_dim}}; } std::string TransposeObj::toString() const diff --git a/src/operators/unary.cc b/src/operators/unary.cc index 3daad361..12d581c0 100644 --- a/src/operators/unary.cc +++ b/src/operators/unary.cc @@ -39,7 +39,9 @@ namespace infini // TODO：返回经过 clip 操作后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13 // =================================== 作业 =================================== - return std::nullopt; + // clip 按元素截断数值，不改变形状，输出形状与输入一致； + // 返回类型是 optional>，返回 {{inputs[0]->getDims()}}。 + return {{inputs[0]->getDims()}}; } std::string ClipObj::toString() const @@ -66,7 +68,11 @@ namespace infini // REF_FILE: src/core/operator.cc // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 // =================================== 作业 =================================== - return {}; + // cast 只有一个输出，其数据类型由 castType 决定； + // 现成的 getOutputDataType() 已实现 CastType -> DataType 的映射， + // 这里直接 return {getOutputDataType()} 即可（否则会继承默认实现 + // "输出类型 = 输入类型"，导致测例里 getOutDType() 断言失败）。 + return {getOutputDataType()}; } optional> CastObj::inferShape(const TensorVec &inputs) @@ -75,7 +81,8 @@ namespace infini // TODO：返回经过 cast 操作后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 // =================================== 作业 =================================== - return std::nullopt; + // cast 只改变数据类型，形状与输入一致，返回 {{inputs[0]->getDims()}}。 + return {{inputs[0]->getDims()}}; } std::string CastObj::toString() const diff --git a/src/utils/operator_utils.cc b/src/utils/operator_utils.cc index edbd2c82..29846805 100644 --- a/src/utils/operator_utils.cc +++ b/src/utils/operator_utils.cc @@ -1,5 +1,7 @@ #include "utils/operator_utils.h" #include "core/runtime.h" +#include "core/tensor.h" +#include namespace infini { @@ -9,8 +11,24 @@ Shape infer_broadcast(const Shape &A, const Shape &B) { // TODO：对 A 和 B 进行双向广播，返回广播后的形状。 // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md // =================================== 作业 =================================== - - return {}; + // - 从后往前对齐两个 shape（可用反向迭代器 rbegin/rend）， + // 较短的 shape 缺失的维度视为 1； + // - 每维规则：两值相等、或其一为 1 时可广播，结果取较大者； + // 否则断言失败； + // - 因为是从后往前生成的，记得 std::reverse 回来 + // （需要 #include ）。 + Shape output_shape; + auto itA = A.rbegin(); + auto itB = B.rbegin(); + while(itA != A.rend() || itB != B.rend()) { + int dimA = (itA != A.rend()) ? *itA++ : 1; + int dimB = (itB != B.rend()) ? *itB++ : 1; + IT_ASSERT(dimA == dimB || dimA == 1 || dimB == 1, + "Broadcast failed: " + vecToString(A) + " vs " + vecToString(B)); + output_shape.emplace_back(std::max(dimA, dimB)); + } + std::reverse(output_shape.begin(), output_shape.end()); + return output_shape; } int get_real_axis(const int &axis, const int &rank) {