4. Metal资源管理:MTLBuffer、MTLTexture、MTLHeap的创建与生命周期

聊到Metal加速,绕不开资源管理这一关。你想想看,模型推理说白了就是算力加数据搬运,数据放哪儿、怎么放、什么时候释放,直接决定了你的App是丝般顺滑还是卡成PPT。很多刚接触Metal的同学最爱犯的一个毛病——把GPU内存当CPU内存使,创建对象跟不要钱似的,结果呢?要么爆显存,要么内存泄漏跑一天流量全耗在debug上。

我个人习惯,把Metal资源管理拆成三块:缓冲区(MTLBuffer)、纹理(MTLTexture)和堆(MTLHeap)。它们就像工具箱里的扳手、螺丝刀和收纳盒,各司其职。但在实际项目中,特别是做端侧推理时,怎么用好这三板斧可有不少讲究。

Metal 资源体系与生命周期 MTLDevice MTLBuffer MTLTexture MTLHeap 适合存储 模型权重、常量、统一数据 适合存储 中间特征图、输入/输出图像 管理子资源 存放 Buffer / Texture 实例 创建 → 配置 → 提交 创建 → 设置 → 采样 创建 → 分配 → 重用 资源回收:setPurgeableState / Heap reset 关键:资源类型决定存储位置和读写方式;Heap 是减少分配开销的利器 推理阶段应提前预分配资源,避免在渲染循环中频繁创建/销毁

4.1 MTLBuffer:GPU的「大通铺」

说白了,MTLBuffer 就是 GPU 能访问的一块连续内存块。它不像纹理那样有格式、维度、Mipmap 的概念,就是一块扁平的、按字节对齐的存储空间。模型权重、矩阵运算的中间数据、甚至反卷积的系数,多半都存在 Buffer 里。

我刚开始做项目的时候,以为 Buffer 不就是 malloc 塞给 GPU 吗?太天真了。这里最关键的是 MTLResourceOptions 这个枚举,它决定了你数据在 CPU 和 GPU 之间的可见性以及存储位置。

选项 存储位置 CPU 访问 适合场景
MTLResourceStorageModeShared 统一内存(UMA) 直接读写 频繁 CPU/GPU 交互的数据
MTLResourceStorageModePrivate GPU 专属显存 不可直接访问 推理权重、固定中间结果
MTLResourceStorageModeManaged CPU + GPU 副本(macOS) 需同步 macOS 平台共享数据

碰见一个常见的误区:有人图方便,所有 Buffer 都用 Shared 模式。这在模型小的时候没问题,一旦模型大了,每次 GPU 计算完 CPU 还得傻等数据同步,反而拖慢。我的经验是:模型权重和运算中间结果用 Private,只在推理开始前通过 Blit 编码器拷贝进去,平时 CPU 压根不去碰它。

// 创建 Private 模式 Buffer,存放推理权重
id<MTLBuffer> weightBuffer = [device newBufferWithLength:weightSize
                                                 options:MTLResourceStorageModePrivate];

// 用 Blit 命令从 Shared 暂存区拷贝到 Private Buffer
id<MTLCommandBuffer> cmdBuf = [commandQueue commandBuffer];
id<MTLBlitCommandEncoder> blitEncoder = [cmdBuf blitCommandEncoder];
[blitEncoder copyFromBuffer:stagingBuffer
               sourceOffset:0
                   toBuffer:weightBuffer
          destinationOffset:0
                       size:weightSize];
[blitEncoder endEncoding];
[cmdBuf commit];

还有个小细节:Buffer 的对齐要求。通常 GPU 喜欢 256 字节对齐,特别是当你准备把这个 Buffer 当作参数传给 Kernel 的时候。如果你的大小为 100 字节,建议按 256 对齐来分配,不要卡着边儿分配,否则 Shader 里读地址很容易越界。

4.2 MTLTexture:带「身份证」的存储区

Texture 和 Buffer 最大的区别是:Texture 知道自己的形状(宽度、高度、深度)、像素格式(MTLPixelFormat)、还有 Mipmap 层数。你在 Metal Shader 里用 texture2d 读写,本质上是 GPU 硬件层面帮你做坐标寻址、过滤、边界处理等。这对图像类任务简直是开挂。

在推理加速里,卷积层的输入输出特征图最适合用 Texture 表示。举个例子,一个 1x3x224x224 的输入,我习惯创建一个 MTLPixelFormatRGBA16Float 的 2D 纹理,正好对应 4 通道。

MTLTextureDescriptor *desc = [[MTLTextureDescriptor alloc] init];
desc.textureType = MTLTextureType2D;
desc.pixelFormat = MTLPixelFormatRGBA16Float;
desc.width = 224;
desc.height = 224;
desc.usage = MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite;

// 这里把整个纹理放到 Private Storage
desc.storageMode = MTLStorageModePrivate;

id<MTLTexture> inputTexture = [device newTextureWithDescriptor:desc];

使用纹理有一个非常容易被忽略的陷阱:Texture 的 ready 状态。新创建的纹理,其内容是不确定的,你必须显式写入数据(比如用 Blit 或 Shader Write),然后再设一个 barrier,保证后续的读取能看到写好的内容。这属于低级错误,但我自己就翻过车——debug 了半天,发现 Texture 里全是随机花花绿绿的数据。

⚠️ 警告:纹理复用与生命周期
千万不要在 MTLCommandBuffer 正在使用 Texture 时,对其进行修改或释放。Metal 使用引用计数管理纹理,如果你 retain 没做好,当前 GPU 计算还没结束,关联的 Texture 又被回收了,那么恭喜你——接下来的 EXC_BAD_ACCESS 跑不了。

如果你创建一个纹理专门用于 shader write,记得把 usage 里加上 MTLTextureUsageShaderWrite。默认只开了读权限。万一你后面想用它存中间结果,却没有写权限,那就只能干瞪眼了。

4.3 MTLHeap:高性能的「豆腐块」管理者

这个是我最想聊的部分。在实际项目中,如果你一个一个地分配 Buffer 和 Texture,Alloc 开销会非常吓人。尤其在做实时推理时,每一帧都要创建输出纹理、临时 Buffer,频繁分配和释放 Direct Memory 会给系统带来很大压力。

MTLHeap 就是用来解决这个问题的。你可以提前申请一大块显存,然后在上面切出小块的 Buffer 或 Texture 来用。它不会产生新的 Allocation,只是从你预留的空间里划出一块。

MTLHeapDescriptor *heapDesc = [[MTLHeapDescriptor alloc] init];
heapDesc.size = 64 * 1024 * 1024; // 64MB
heapDesc.storageMode = MTLStorageModePrivate;
// iOS 建议使用 MTLHazardTrackingModeTracked
heapDesc.hazardTrackingMode = MTLHazardTrackingModeTracked;

id<MTLHeap> heap = [device newHeapWithDescriptor:heapDesc];

// 从 Heap 里分配 Texture 和 Buffer
MTLTextureDescriptor *texDesc = [MTLTextureDescriptor texture2DDescriptorWithPixelFormat:MTLPixelFormatRGBA16Float
                                                                                   width:224
                                                                                  height:224
                                                                               mipmapped:NO];
texDesc.usage = MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite;

id<MTLTexture> heapTex = [heap newTextureWithDescriptor:texDesc];
id<MTLBuffer> heapBuf = [heap newBufferWithLength:4096 options:0];

要注意的是,Heap 上的资源不能设置 storageMode,因为 Heap 创建时已经统一决定了存储模式。另外,从 Heap 分配的资源生命周期跟 Heap 绑定,Heap 释放了,所有子资源自动作废。所以你必须在推理阶段用完后,及时做 heap reset 或让 Heap 进入 purgeable 状态,才能把内存还给 GPU。

对比项 单独分配资源 从 Heap 分配
分配开销 较高(每次走内核分配) 低(仅子分配)
内存碎片 容易产生 Heap 内部可复用
资源回收 逐个 release reset Heap 批量回收
适用场景 少量、长期持有资源 大量临时资源(推理中间结果)

我曾在做 YOLO 推理优化时,把所有卷积层的输出纹理统一放在一个 Heap 里。每层计算完,把下一层需要的纹理就从 Heap 里再申请出来。等一整个推理流程结束,直接 reset Heap,下次推理时又能复用。开销主要就是前期的一次 Heap 声明,后面全是零内存分配。

💡 建议:
如果你的模型推理函数被频繁调用,建议用 MTLHeap + setPurgeableState 手动管理,而不是依赖 ARC 或 Metal 的自动回收。自动回收在天亮之前是好的,但推理对时序敏感,一旦触发 Page Out 就掉帧。

4.4 资源生命周期管理最佳实践

结合前面讲的三类资源,我把推理场景下的资源管理套路总结成几步:

  1. 预分配阶段:模型加载时,把权重、常量用 Private 模式的 Buffer 一次性写进去。这部分资源生命期跟推理会话一致。
  2. Heap 预热阶段:根据模型最大吞吐量,计算每层特征图的最大尺寸总和,开一个足够大的 Heap。所有中间纹理、临时 Buffer 都从这里拿。
  3. 推理循环:从 Heap 获取纹理/Buffer → 执行命令 → 处理完后直接丢弃引用。Heap 不释放,资源就一直在被循环利用。
  4. 清理阶段:整个推理会话结束时,手动 [heap reset],然后释放 Heap 本身。

还有一个小坑:Hazard Tracking。Metal 3 引入了 MTLHazardTrackingModeTracked。要是你开了 Tracked,Metal 自动会加 barrier 和 fence,防止读写冲突。好处是省心,坏处是可能加了一些你没打算要的同步点。追求极致性能时,我会关掉这个选项,自己用 fence 来控制同步,但前提是你得清楚每条命令的数据依赖——这个我们后续章节会展开。

到最后,GPU 内存的本质是你要学会「控制节奏」。Buffer 像大通铺,什么东西都能放;Texture 像带格式的专业板房,给图像用最合适;Heap 就像整栋宿舍楼,让你统一分配不用来回敲内核。组合起来用,你的推理跑起来才会又快又稳。