基于WebGPU实现的热力图

背景

笔者最近在重构一个地理数据可视化工具，其中渲染器由之前的WebGL改为WebGPU实现并抽离出来形成了一个单独的项目，地址尚处于开发阶段。目前初步支持了散点和轨迹的渲染，二者实现较为简单。但对于热力数据渲染的实现，WebGPU和WebGL确有很大的不同

渲染热力图原理

简介

热力图是利用利用不同的颜色（通常按照红、黄、绿、蓝）表征某个事件发生概率高低的图表，常用于地图应用中，如openlayers 官网提供的利用热力地图表示全球地震发生概率的示例

实现简介

热力值

热力值表示事件发生概率的高低，等于事件在某个像素上发生的概率值取值为[0, 1]。渲染热力图主要就是计算各个像素上的热力值，然后通过颜色插值获取到实际的 rgb 颜色。

热力点

热力点包含两个要素，一个是点的像素坐标，另一个是点的像素半径。热力图应用中，后端往往以二维点为单位组织数据（比如经纬度点）。我们拿到后端的点数据后需要先将点的实际坐标经过一系列坐标变换转换到屏幕空间上，再加上提前定义好的像素半径，就得到了各个热力点数据。

通过热力点数据计算各个像素的热力值

这一步是实现热力图渲染的关键，我们在拿到热力点数据后，可以通过各个热力点的像素坐标和像素半径计算出各个热力点覆盖的所有像素，并计算出单个像素上的热力值之和。这个步骤的实现基于以下几个约定：

热力点中，圆心处中的热力值最高（设为1），圆周处的热力值最低（设定为0）。单个热力点内的各个像素上的热力值以像素到圆心距离成正比
因为单个像素可能被多个热力点覆盖，像素在各个热力点上的热力值之和可能大于1，所以在颜色插值前允许像素上的热力值大于1。后续在颜色插值时对各个像素的热力值做0到1的归一化

该步骤的大致实现可以细分为以下几步：

准备与屏幕分辨率匹配的内存空间，比如js 的 canvas 或者 webgl 的 texture 或者 webgpu 的 buffer，用于存放每个像素的热力值，假定为 pixelHeatValueArray
遍历所有热力点，针对每个热力点计算该热力点覆盖的所有像素
遍历步骤2中的像素，计算该像素到当前热力点的距离，根据该距离得到当前热力点在该像素上的热力值
将步骤3得到的热力值累加到pixelHeatValueArray。重复步骤3直到当前热力点覆盖的像素点遍历完毕
重复步骤2直到所有热力点遍历完毕，最后输出 pixelHeatValueArray。

这一步可以利用 js 在 CPU 中完成，也可以利用 webgl 或者 webgpu 等图形库在 GPU 中并行完成。

热力值到 RGB 颜色的插值

上一小节中得到的各个像素的热力值只是一个未做归一化处理的浮点数，还无法作为颜色输出到画布上。最后我们需要根据最大热力值（可以是用户设置的固定值或者遍历 pixelHeatValueArray 数组后得到），各个像素上的热力值以及设置的颜色数组，插值得到各个像素上的 RGB 颜色。并输出到画布上。这一步同样可以利用 js 在 CPU 中完成，也可以利用 webgl 或者 webgpu 等图形库在 GPU 中并行完成。

旧版本热力数据渲染WebGL的实现

实现过程

通过上一章节的介绍我们可以知道，渲染热力图主要的两个步骤就是像素点上热力值的计算，以及像素点上热力值到 RGB 颜色的插值。这两步是承前启后的关系，无法并行完成。又因为 webgl 中没有计算着色器的概念，所以在 webgl 实现的版本中，我们利用了两个渲染管线来完成热力图的渲染。

第一个渲染管线中，我们将热力点作为点模型进行离屏渲染，在顶点着色器中完成点的经纬度坐标到裁剪空间的坐标转换，后续由 webgl将点的坐标进一步转换到屏幕空间中。
接着通过光栅化阶段后，webgl 会对各个热力点覆盖的方形区域下的像素运行片元着色器。在片元着色器中我们计算得到各个像素在热力点下的热力值并作为gl_FragColor.r 分量输出，注意该热力值只是覆盖当前像素的某个热力点的下的热力值。
在webgl渲染管线的混合过程中，我们设置片元着色器输出到颜色缓存时目标颜色和源颜色的各个分量进行系数为1的叠加，这样在混合过程中我们就完成了不同热力点在同一个像素上热力值的累加操作。
第一个渲染管线输出的纹理作为第二个渲染管线的输入纹理。第二个渲染管线渲染一个和屏幕同样大小的矩形区域，在其片元着色器中读取各个像素对应的纹理颜色，通过读取红色分量获取像素对应的热力值。然后通过颜色插值得到最终的RGB 颜色输出到屏幕上。

缺点

在第一个渲染管线中，热力值累加到片元着色器输出颜色的红色分量，该分量的数值最终输出到纹理后只有8bit的长度，精度太低
整个过程无法获取第一个渲染管线输出的最大热力值。无法对后续的显示效果进行优化（比如渲染上百万个热力点时，实际最大的热力值可能几万，如果利用实际最大热力值对像素热力值做归一化，可能导致页面看着全是红色），从而需要用户通过页面交互手动调整最大热力值，得到优化渲染效果的目的

WebGPU实现热力数据渲染

我们在WebGPU 版本中最大的改进就是利用compute shader完成像素点上热力值的计算以及最大热力值的获取。由于最大热力值的计算是在 GPU 中完成的，所以不再需要将所有像素的热力值从显存拷贝到内存中，减少了性能开销。按照逻辑上的先后顺序，我们将实现步骤大致分为计算像素热力值、计算最大热力值以及 RGB 颜色插值三个主要步骤。其中前两个步骤由通过两个compute pass计算通道完成，RGBA颜色插值由一个render pass渲染通道完成

计算像素热力值

该步骤其实就是完成webgl热力图渲染版本中第一个渲染管线的功能，即拿到各个热力点的坐标后计算屏幕上所有像素点上的热力值。输入包括

一个保存了各个热力点经纬度坐标（以地图热力图为例）的storage buffer input（WebGPU中的一类buffer，具有空间大在compute shader中可读写的功能）
一个用于保存各个像素点热力值的storage buffer output
一个保存了屏幕分辨率的二维浮点数向量resolution
热力点的像素半径radius
包含webgpu在各个维度(X/Y/Z)上调用的workgroup(工作组)的数量的三维数组grid。如果compute shader代码中定义的workgroup_size为(1, 1, 1)即每个工作组内部只有一个workItem工作项，那么就可以用grid表示compute shader实际执行的次数
将热力点的实际坐标变换到裁剪空间下的投影矩阵和视图矩阵

实现步骤为：

1. 坐标变换

需要注意的是在光栅化的顶点着色器中，输出到顶点坐标是位于裁剪空间内的点的XYZ坐标取值范围为[-w, w], [-w, w]和[0, w]。后续GPU会自动将裁剪空间内的点的坐标转换到NDC标准设备空间（xy的取值范围为[-1, 1]）然后乘以设备的分辨率转换到二维屏幕空间上

我们在compute shader中将热力点坐标转换为屏幕像素坐标时，先要用投影矩阵和视图矩阵乘以点的原始坐标得到热力点在裁剪空间下的坐标。然后将坐标的xyz分量除以齐次坐标的第四个分量齐次分量w，这样就得到NDC空间下的坐标，最后加上二维向量 vec2(1, 1)后再乘以resolution，就得到热力点在当前屏幕上的像素坐标了。这是后续计算各个像素点热力值的前提。 &emps;&emps;具体的坐标转换代码为

rust 复制代码

let pointInClip = projectionMatrix * viewMatrix * vec4f(input[index], 0, 1);
//将裁剪空间下的 point 坐标转换为屏幕空间的像素坐标，其中像素坐标的原点位于屏幕的左下方
let pointInScreen = (pointInClip.xy / pointInClip.w + vec2f(1, 1)) / 2 * resolution ;

其中input为存放了所有热力点原始二维坐标的storage buffer。index为当前compute shader计算的热力点的index。compute shader输入为三维向量global_invocation_id，因为webgpu的workgroup是三维布局，所以根据global_invocation_id和输入的grid可以计算出global_invocation_id对应的热力点的index。

2. 计算像素点上的热力值

在得到热力点的屏幕像素坐标之后，我们在这一步中根据输入的热力点像素半径遍历以热力点为中心的2*radius边长的正方形区域。依次计算热力点在各个像素上的热力值，然后将热力值累加到output buffer中。具体实现代码如下

rust 复制代码

for(var i = si; i <= ei; i++){
    for(var j = sj; j <= ej; j++){
        let d = pow(pow(f32(i) - pc.x, 2) + pow(f32(j) - pc.y, 2), 0.5);
        var h = step(d / radius, 1) * (1 - d / radius);
        h = pow(h, 1.5);
        let outIdx = u32(j) * u32(resolution.x) + u32(i);
        let v = u32(step(0, h) * h * prec);
        atomicAdd(&output[outIdx], v);
    }
}

变量d为像素点到热力点中心的距离单位为像素，通过d与半径radius的比值求出像素点的热力值。outIdx为像素点在output buffer中的索引。

因为一个compute shader只负责处理一个热力点，在compute shader像output buffer中累加某个像素点的热力值时，可能存在处理其它热力点的compute shader也在往同一个像素点累加热力值，而常规的累加操作可以分为从output buffer中读入热力值、热力值累加以及往output buffer中写入累加后的热力值三个操作。所以多个compute shader并行处理的情况下可能会出现被多个热力点覆盖的像素点上的热力值计算错误的问题。此处我们使用了webgpu的原子操作，通过atomicAdd实现累加的原子操作。

3. 计算所有像素的最大热力值

前两步中通过一个compute pass我们得到了所有像素点上的热力值，为了对热力值进行0到1的归一化，我们还需要得到所有像素点上的最大热力值。这个在cpu中很好实现，只用遍历一般output buffer即可。但output buffer位于显存中，为了完成遍历需要讲output 映射到cpu可以访问的内存中，这个过程我们测试过比较耗时（1080分辨率的output buffer在2018款mbp中大概需要0.3秒远大于一帧渲染的耗时）

所以我们通过第二个compute pass使用GPU并行计算来获取最大热力值。具体步骤为

调用resolution.y个compute shader，每个compute shader统计一行像素中的最大热力值
计算完一行像素的最大热力值之后，通过atomicMax原子操作将最大值与之前的最大热力值比较
输出最终的最大热力值

4. RGBA颜色插值

颜色插值的实现和webgl版本热力图差不多。准备一个渲染管线，创建一个屏幕大小的矩形几何体，在fragment shader中通过shader的输入参数position求得当前着色器处理的像素在output buffer中的索引，然后对热力值进行归一化，再根据输入的颜色数组对热力值进行插值，最后输出颜色。下图是一万个随机点半径为15pixel的热力图