用图形GPGPU做物理仿真（2）·双调排序（上）

Doris232 · 发表于 2022-12-24 15:57

从这篇文章开始，我们将正式进入到GPGPU的开发当中。想来想去我觉得从一个排序算法开始时比较好的，这样一来可以先把基础的架构部分梳理清楚，并且不涉及到任何数值相关的内容。所以，我会围绕着双调排序算法的GPGPU实现，介绍使用的一些接口，通过组合这些接口简化GPGPU开发的复杂性，为叠加更复杂的算法做准备。
GPU排序是我们后续必然会用到的模块，调研下来一般来说会有两种方法，双调排序BitonicSort，和RadixSort，这里介绍的是双调排序，可以直接在原数组基础上进行排序，不需要额外的GPU显存分配。排序方法将处理 float2 类型的数组，x 设置成用于排序比较的key，y 则是粒子的index。因此可以用距离相机的距离作为key，从而排序透明渲染的粒子，也可以将空间网格的hash值作为key，以获得同一网格当中的粒子用于计算例如SPH的Kernel值。
本文的代码位于：
<hr/>算法原理

双调排序的原理使得他非常适合并行实现，算法可以参考：
Batcher定理保证将任意一个长为2n的双调序列A分为等长的两半X和Y，将X中的元素与Y中的元素一一按原序比较，即a与a[i+n] (i < n)比较，将较大者放入MAX序列，较小者放入MIN序列。则得到的MAX和MIN序列仍然是双调序列，并且MAX序列中的任意一个元素不小于MIN序列中的任意一个元素。具体算法实现可以参考：
在仓库当中的文档有一个动画，比较直观得展示了双调排序的过程：
对于一串数列，先左右两个组合排序成双调数列，然后逐级归并排序，最终就可以得到单调的数列。

对于这样一个算法如果直接转换成计算着色器，一个kernel就可以完成，但如果要考虑优化，则需要拆分成三个kernel。在开发GPGPU时最主要的优化点是利用显式控制GPU内存的使用，和CPU上程序员只能通过对齐来利用多级缓存不同，GPU上可以显式设定数据分配在哪一块内存区域上。
因此对于比较短的数列，可以在一个threadgroup中分配共享的内存，避免kernel直接从device memory当中直接存取数据带来的延时。所以，在实现中使用3个kernel：

pre-sort：对数据按照512为一组进行排序，每一块都在一个threadgroup当中进行，共享内存

step-sort：对于超过512的数组，对于大组之间进行排序

inner-sort：对于大组排序后数组，再对512块内的数据进行排序

pre-sort只会执行一次，step-sort 和 inner-sort 会根据数组长度迭代执行。
<hr/>缓存 Shader 和 PSO

对于任何图形API的程序，都会涉及到编译 Shader，构造 PSO 等管线构建的过程。所以你写完shader之后，还需要写一大串的管线搭建代码，并且把数据绑定到管线上。因为这些东西所以写起来会比CUDA复杂很多。Shader和PSO都需要根据hash缓存下来，如果每一帧都重新构建会严重影响性能。对于计算管线来说，管线的状态比较简单，因此直接根据编译出的Shader地址作为Shader的HashValue，同时可以计算PSO的HashValue。在Swift当中 NSObject对象自带hash的属性，直接拿过来用就好了：
func requestShaderModule(_ shaderPass: ShaderPass, _ macroInfo: ShaderMacroCollection) -> [MTLFunction] {
var functions: [MTLFunction] = []
for shader in shaderPass._shaders {
      var hasher = Hasher()
      shader.hash(into: &hasher)
      macroInfo.hash(into: &hasher)
      hasher.combine(shaderPass._library.hash)
      let hash = hasher.finalize()
      let cacheFunction = shader_modules[hash]
      if cacheFunction == nil {
         let function = shaderPass.createProgram(shader, macroInfo)
         if function != nil {
            shader_modules[hash] = function!
            functions.append(function!)
         }
      } else {
         functions.append(cacheFunction!)
      }
}
return functions
}

func requestComputePipeline(_ pipelineDescriptor: MTLComputePipelineDescriptor) -> ComputePipelineState {
let hash = pipelineDescriptor.hash
var pipelineState = compute_pipelines[hash]
if pipelineState == nil {
      pipelineState = ComputePipelineState(device, pipelineDescriptor)
      compute_pipelines[hash] = pipelineState
}

return pipelineState!
}如果你使用的是C++，那就需要自己去计算hashValue，例如对于WebGPU可以采用：
template <>
struct hash<wgpu::ComputePipelineDescriptor> {
std::size_t operator()(const wgpu::ComputePipelineDescriptor &descriptor) const {
      std::size_t result = 0;

      hash_combine(result, descriptor.layout.Get());  // internal address
      hash_combine(result, descriptor.compute);

      return result;
}
};
<hr/>反射

缓存 Shader 和 PSO 之后，我们只需要关算子Kernel和数据这两项就可以了。在CUDA当中，可以像调用函数一样直接编写如下的代码：
matMulKernel << < gridSize, blockSize >> >(A, B, C);但对于图形API则需要通过binding才能绑定数据。特别的还需要制定binding的index，Metal的binding其实是比较简单的，有点偏向传统的OpenGL风格，但如果是DX和Vulkan，绑定需要更加繁琐的代码。所以我们需要一种机制将数据和算子解耦，将绑定的过程自动执行，这样就可以简化大量的代码。
实际上，在Shader编写时我们就指定了每一种数据的类型和index，如果能够获取这些参数，并且查找这些参数对应的数据，就获得了绑定所需的所有信息。因此，我们需要反射Shader的信息，基于这些信息自动执行绑定。在 Metal 当中可以通过 MTLComputePipelineReflection 对象获取管线的信息。对于其他API，普遍会使用SPIRV的工具链得到这些数据。
do {
_handle = try device.makeComputePipelineState(descriptor: descriptor,
         options: MTLPipelineOption.argumentInfo, reflection: &_reflection)
} catch let error {
fatalError(error.localizedDescription)
}为了应用反射出来的数据，我们将所有GPU内存管理都在一个集中的仓库中进行，那就是ShaderData，这个类型可以方便地存储Buffer，Texture等数据，也方便我们对这些数据进行构造和释放。例如可以有这样一个泛型函数：
public func setData<T>(_ property: String, _ data: T) {
let value = _shaderBuffers.first { (key: String, value: BufferView) in
      key == property
}
if value == nil {
      _shaderBuffers[property] = BufferView(device: engine.device, array: [data])
} else {
      value!.value.assign(data)
}
}将数据集中，就可以通过在这些数据集当中搜索和Shader参数同名的数据，实现绑定的操作：
func bindData(_ commandEncoder: MTLComputeCommandEncoder,
            _ reflectionUniforms: [ReflectionUniform],
            _ resourceCache: ResourceCache) {
for uniform in reflectionUniforms {
      switch uniform.bindingType {
      case .buffer:
         if let buffer = _shaderBuffers[uniform.name] {
            commandEncoder.setBuffer(buffer.buffer, offset: 0, index: uniform.location)
         }
         if let buffer = _shaderDynamicBuffers[engine.currentBufferIndex][uniform.name] {
            commandEncoder.setBuffer(buffer.buffer, offset: 0, index: uniform.location)
         }
         if let bufferFunctor = _shaderBufferFunctors[uniform.name] {
            commandEncoder.setBuffer(bufferFunctor().buffer, offset: 0, index: uniform.location)
         }
         break
      case .texture:
         if let image = _imageViews[uniform.name] {
            commandEncoder.setTexture(image, index: uniform.location)
         }
         break
      case .sampler:
         if let sampler = _samplers[uniform.name] {
            commandEncoder.setSamplerState(resourceCache.requestSamplers(sampler), index: uniform.location)
         }
         break
      default:
         break
      }
}
}<hr/>ComputePass

有了这些工具我们就可以整合到ComputePass当中，和用于渲染的RenderPass不同，我在设计这个类型的时候只是希望在调用Dispatch的时候能更加容易，自动创建PSO，自动用反射信息binding，自动用最优的方式设置dispatchGroup的数量等等。因此整个封装是非常薄的，从而使得计算着色器的使用可以更加灵活，不会被封装所限制。所以我们可以封装这样一个函数：
open func compute(commandEncoder: MTLComputeCommandEncoder, label: String = &#34;&#34;) {
      commandEncoder.pushDebugGroup(label)
      if let resourceCache = resourceCache {
         let compileMacros = ShaderMacroCollection()
         for shaderData in data {
            ShaderMacroCollection.unionCollection(compileMacros, shaderData._macroCollection, compileMacros)
         }

         for shaderPass in shader {
            _pipelineDescriptor.computeFunction = resourceCache.requestShaderModule(shaderPass, compileMacros)[0]
            let pipelineState = resourceCache.requestComputePipeline(_pipelineDescriptor)
            for shaderData in data {
                  shaderData.bindData(commandEncoder, pipelineState.uniformBlock, resourceCache)
            }
            commandEncoder.setComputePipelineState(pipelineState.handle)

            let nWidth = min(threadsPerGridX, pipelineState.handle.threadExecutionWidth)
            let nHeight = min(threadsPerGridY, pipelineState.handle.maxTotalThreadsPerThreadgroup / nWidth)
            commandEncoder.dispatchThreads(MTLSize(width: threadsPerGridX, height: threadsPerGridY, depth: threadsPerGridZ),
                     threadsPerThreadgroup: MTLSize(width: nWidth, height: nHeight, depth: 1))
         }
      }
      commandEncoder.popDebugGroup()
}这里只需要说两点，上一篇文章介绍 Debug 工具时，我没有说的一点就是，其实Debug工具好不好用，和代码里面是不是充分利用 label 标签和 DebugGroup 有关，如果尽可能使用这两个特性对管线的流程进行标记，那么抓帧之后能够更容易知道当前这一步在做什么操作。
另外一点，Metal有一点和DX，Vulkan都不同，一般来说我们分配的thread数量会大于需要的，比如图像处理的大小是934x283，那么thread很难完全覆盖，对于边界情况要在shader里做处理。但Metal的 dispatchThreads 接口简化了这一操作，直接dispatch 934x283x1 个thread，对于边界的情形，管线会自动做处理。进一步还可以使用 threadExecutionWidth 属性计算出当前设备上并行执行的thread数量，用这一参数作为 threadsPerThreadgroup。这样一来用户就没必要纠结threadgroup的数量和threadPerThreadGroup的数量到底要怎么设置，只需要关注thread的总量就好了，更加便捷。
但是有时候我们需要更加细致地控制所需要的threadgroup的数量和threadPerThreadGroup的数量，就可以使用如下的接口：
commandEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)在这个例子当中还需要 indirect buffer 去dispatch，因为后面会讲到我们会在 GPU 当中做粒子的发射器，我们需要用一个 atomic buffer 作为计数器，而我们对粒子进行排序时，所需要的threadgroup数量也是基于粒子总量进行计算的，所以使用 indirect buffer 触发dispatch会更方便：
commandEncoder.dispatchThreadgroups(indirectBuffer: indirectBuffer,
                                 indirectBufferOffset: 0, threadsPerThreadgroup: threadsPerThreadgroup)有了这样的封装，我们只需要做一些基础的配置，指定computePass使用的shader和shaderData，然后调用compute接口就可以很容易调用计算着色器管线。这样一来我们的注意力只需要集中在算法本身就可以了。
<hr/>下一篇文章会回到双调排序算法本身，介绍着色器的编写以及多个computePass的组织。

rustum · 发表于 2022-12-24 16:03

网上有CUDA版本，也有ComputeShader版本，OpenCL版本的没有关注过；GPU排序，给GPU用够快，如果再给CPU用，还真不知道，有测试过的告知下。

		自动登录	找回密码
密码			立即注册

用图形GPGPU做物理仿真（2）·双调排序（上）

本帖子中包含更多资源