2D Jacobi Stencil

原始题目：LeetGPU - 2D Jacobi Stencil

题目描述

给定一个 32 位浮点数的二维网格，应用一次五点 Jacobi 模版迭代：输出中每个内部单元的值设为其在输入网格中的四个基本方向邻居（上、下、左、右）的平均值。边界单元（第一行/最后一行和第一列/最后一列）从输入原样复制到输出。

对于内部单元 $(i, j)$（$0 < i < rows-1,\ 0 < j < cols-1$）：

$$output[i, j] = \frac{input[i-1, j] + input[i+1, j] + input[i, j-1] + input[i, j+1]}{4}$$

对于边界单元：$output[i, j] = input[i, j]$。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 中。
• 仅从 input 读取，仅向 output 写入。

示例

Input (4×4): [[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]]
Interior cell (1,1): (2+10+5+7)/4 = 6.0
Interior cell (1,2): (3+11+6+8)/4 = 7.0
Interior cell (2,1): (6+14+9+11)/4 = 10.0
Interior cell (2,2): (7+15+10+12)/4 = 11.0
Output: 边界原样, 内部 [[6,7],[10,11]]

约束条件

• $1 \le rows, cols \le 16{,}384$。
• 输入值范围 $[-100, 100]$，均为 32 位浮点数。
• 性能测试在 $rows = cols = 8{,}192$ 的规模下进行。

解题思路

Jacobi 模版是典型的邻域访问（Stencil）模式。每个输出元素需要读取上下左右四个邻居，相邻输出之间大量重叠。核心优化是使用共享内存分块并包含 halo 区域（边界外的一圈），让线程块内的线程协作加载数据。边界处理需要条件判断以避免越界访问。模版计算是内存带宽受限的——每次迭代有 4 次读取和 1 次写入，算术强度极低。欢迎在 GitHub Discussions 分享你的解法。