wgpu/src/utils/mip-generator.ts

import code from './mip-shader.wgsl' with { type: 'text' }

const mip = (n: number) => Math.max(1, n >>> 1)

export class MipGenerator {
  _device: GPUDevice
  _sampler: GPUSampler
  _pipelines: Record<string, GPURenderPipeline>

  _shader?: GPUShaderModule
  _bindGroupLayout?: GPUBindGroupLayout
  _pipelineLayout?: GPUPipelineLayout

  constructor(device: GPUDevice) {
    this._device = device
    this._sampler = device.createSampler({ minFilter: 'linear' })
    this._pipelines = {}
  }

  _getShader() {
    if (!this._shader) {
      this._shader = this._device.createShaderModule({
        code
      })
    }

    return this._shader
  }

  _getBindGroupLayout() {
    if (!this._bindGroupLayout) {
      this._bindGroupLayout = this._device.createBindGroupLayout({
        entries: [{
          binding: 0,
          visibility: GPUShaderStage.FRAGMENT,
          sampler: {}
        }, {
          binding: 1,
          visibility: GPUShaderStage.FRAGMENT,
          texture: {}
        }]
      })
    }

    return this._bindGroupLayout
  }

  _getPipelineLayout() {
    if (!this._pipelineLayout) {
      this._pipelineLayout = this._device.createPipelineLayout({
        label: 'Mipmap Generator',
        bindGroupLayouts: [this._getBindGroupLayout()],
      })
    }

    return this._pipelineLayout
  }

  getPipeline(format: GPUTextureFormat) {
    let pipeline = this._pipelines[format]

    if (!pipeline) {
      const shader = this._getShader()

      pipeline = this._device.createRenderPipeline({
        layout: this._getPipelineLayout(),
        vertex: {
          module: shader,
          entryPoint: 'vs_main'
        },
        fragment: {
          module: shader,
          entryPoint: 'fs_main',
          targets: [{ format }]
        }
      })

      this._pipelines[format] = pipeline
    }

    return pipeline
  }

  generateMipmap(texture: GPUTexture, descriptor: GPUTextureDescriptor) {
    const pipeline = this.getPipeline(descriptor.format)

    if (descriptor.dimension !== '2d') {
      throw new Error('Generating mipmaps for anything except 2d is unsupported.')
    }

    let mipTexture = texture
    const { width, height, depthOrArrayLayers } = descriptor.size as GPUExtent3DDict

    const renderToSource = descriptor.usage & GPUTextureUsage.RENDER_ATTACHMENT

    if (!renderToSource) {
      mipTexture = this._device.createTexture({
        size: {
          width: mip(width),
          height: mip(height),
          depthOrArrayLayers
        },
        format: descriptor.format,
        usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_SRC | GPUTextureUsage.RENDER_ATTACHMENT,
        mipLevelCount: descriptor.mipLevelCount - 1
      })
    }

    const encoder = this._device.createCommandEncoder({})

    for (let layer = 0; layer < depthOrArrayLayers; ++layer) {
      let srcView = texture.createView({
        baseMipLevel: 0,
        mipLevelCount: 1,
        dimension: '2d',
        baseArrayLayer: layer,
        arrayLayerCount: 1
      })

      let dstMipLevel = renderToSource ? 1 : 0
      for (let i = 1; i < descriptor.mipLevelCount; ++i) {
        const dstView = mipTexture.createView({
          baseMipLevel: dstMipLevel++,
          mipLevelCount: 1,
          dimension: '2d',
          baseArrayLayer: layer,
          arrayLayerCount: 1
        })

        const passEncoder = encoder.beginRenderPass({
          colorAttachments: [{
            view: dstView,
            loadOp: 'clear',
            storeOp: 'store'
          }]
        })

        const bindGroup = this._device.createBindGroup({
          layout: this._bindGroupLayout,
          entries: [{
            binding: 0,
            resource: this._sampler
          }, {
            binding: 1,
            resource: srcView
          }]
        })

        passEncoder.setPipeline(pipeline)
        passEncoder.setBindGroup(0, bindGroup)
        passEncoder.draw(3, 1, 0, 0)
        passEncoder.end()

        srcView = dstView
      }
    }

    if (!renderToSource) {
      const mipLevelSize = {
        width: mip(width),
        height: mip(height),
        depthOrArrayLayers
      }

      for (let i = 1; i < descriptor.mipLevelCount; ++i) {
        encoder.copyTextureToTexture({
          texture: mipTexture,
          mipLevel: i - 1
        }, {
          texture,
          mipLevel: i,
        }, mipLevelSize)

        mipLevelSize.width = mip(mipLevelSize.width)
        mipLevelSize.height = mip(mipLevelSize.height)
      }
    }

    this._device.queue.submit([encoder.finish()])

    if (!renderToSource) {
      mipTexture.destroy()
    }

    return texture
  }
}