tensorflow - 为什么 TFLite Metal 中线程组 num (work_group_launch_order) 的顺序很重要？

Question

我是 Metal 的初学者，并试图了解TFLite中卷积的 Metal 实现。在阅读了这行代码并找到了所有用法之后，我真的很困惑为什么这很work_group_launch_order重要。这种机制究竟是如何运作的？它与 GPU 线性化 3D 的方式有关threadgroup吗？

[张量流 GitHub]

为简单起见，让我尝试解释一下ConvolutionGenericTFLite 中某个内核（名为 ）的线程调度策略。

对于给定的 3D 螺纹编号t=<tx, ty, tz>和螺纹组形状，TFLite通过、和s=<sx, sy, sz>计算所需螺纹组的数量。这是绝对正常的。n=<nx, ny, nz>nx=ceil(tx/sx)ny=ceil(tx/sy)nz=ceil(tz/sz)

正常情况下，我们可以分派<nx, ny, nz>3维线程组，通过参数gidattribute获取Metal核函数网格中的线程位置[[thread_position_in_grid]]。线程位置可以决定当前线程负责哪个区域。

然而，TFLite 选择了一种奇怪的方式，将线程组调度为 3 维，并通过计算 from和as<nz, nx, ny>获取 Metal 核函数中网格中的线程位置tid3d [[thread_position_in_threadgroup]]group_id [[threadgroup_position_in_grid]]

gid_x = group_id.y * sx + tid3d.x;
gid_y = group_id.z * sy + tid3d.y;
gid_z = group_id.x * sz + tid3d.z

令我惊讶的是，这种策略确实提高了性能（约 10% 的加速）。

有人可以帮我解释一下奇怪的threadgroup调度策略背后的底层机制吗？

标签： tensorflowmetalgputensorflow-lite