WebGPU Texture - Generate Mipmaps Speed Test (compute mipmap gen)

README.md

WebGPU Texture - Generate Mipmaps Speed Test (compute mipmap gen)

view on jsgist

index.css

@import url(https://webgpufundamentals.org/webgpu/resources/webgpu-lesson.css);
html, body {
  margin: 0;       /* remove the default margin          */
  height: 100%;    /* make the html,body fill the page   */
}
#info {
  position: absolute;
  left: 0;
  top: 0;
  background-color: black;
  color: white;
  margin: 0;
  padding: 0.5em;
}

index.html

<canvas></canvas>
<pre id="info"></pre>

index.js

// WebGPU Simple Textured Quad - Import Canvas
// from https://webgpufundamentals.org/webgpu/webgpu-simple-textured-quad-import-canvas.html


// see https://webgpufundamentals.org/webgpu/lessons/webgpu-utils.html#wgpu-matrix
import {mat4} from 'https://webgpufundamentals.org/3rdparty/wgpu-matrix.module.js';
import GUI from 'https://webgpufundamentals.org/3rdparty/muigui-0.x.module.js';
import RollingAverage from 'https://webgpufundamentals.org/webgpu/resources/js/rolling-average.js';

const range = (num, fn) => new Array(num).fill(0).map((_, i) => fn(i));

async function main() {
  const adapter = await navigator.gpu?.requestAdapter();
  const device = await adapter?.requestDevice();
  if (!device) {
    fail('need a browser that supports WebGPU');
    return;
  }
  device.addEventListener('uncapturederror', e => console.error(e.error.message));

  const fpsAverage = new RollingAverage();

  // Get a WebGPU context from the canvas and configure it
  const canvas = document.querySelector('canvas');
  const context = canvas.getContext('webgpu');
  const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
  context.configure({
    device,
    format: presentationFormat,
  });

  const module = device.createShaderModule({
    label: 'our hardcoded textured quad shaders',
    code: `
      struct OurVertexShaderOutput {
        @builtin(position) position: vec4f,
        @location(0) texcoord: vec2f,
      };

      struct Uniforms {
        matrix: mat4x4f,
      };

      @group(0) @binding(2) var<uniform> uni: Uniforms;

      @vertex fn vs(
        @builtin(vertex_index) vertexIndex : u32
      ) -> OurVertexShaderOutput {
        let pos = array(
          vec3f(-0.5, 0.5,-0.5),
          vec3f( 0.5, 0.5,-0.5),
          vec3f(-0.5, 0.5, 0.5),
          vec3f(-0.5, 0.5, 0.5),
          vec3f( 0.5, 0.5,-0.5),
          vec3f( 0.5, 0.5, 0.5),
        );

        var vsOutput: OurVertexShaderOutput;
        let xyz = pos[vertexIndex];
        vsOutput.position = uni.matrix * vec4f(xyz, 1.0);
        vsOutput.texcoord = (xyz.xz + 0.5) * vec2f(1, 50);
        return vsOutput;
      }

      @group(0) @binding(0) var ourSampler: sampler;
      @group(0) @binding(1) var ourTexture: texture_2d<f32>;

      @fragment fn fs(fsInput: OurVertexShaderOutput) -> @location(0) vec4f {
        return textureSample(ourTexture, ourSampler, fsInput.texcoord);
      }
    `,
  });

  const pipeline = device.createRenderPipeline({
    label: 'hardcoded textured quad pipeline',
    layout: 'auto',
    vertex: {
      module,
    },
    fragment: {
      module,
      targets: [{ format: presentationFormat }],
    },
  });

  const numMipLevels = (...sizes) => {
    const maxSize = Math.max(...sizes);
    return 1 + Math.log2(maxSize) | 0;
  };

  function copySourceToTexture(device, texture, source, {flipY} = {}) {
    device.queue.copyExternalImageToTexture(
      { source, flipY, },
      { texture },
      { width: source.width, height: source.height },
    );

    //if (texture.mipLevelCount > 1) {
    //  generateMips(device, texture);
    //}
  }

/*
  const generateMips = (() => {
    let sampler;
    let module;
    const pipelineByFormat = {};

    return function generateMips(device, texture) {
      if (!module) {
        module = device.createShaderModule({
          label: 'textured quad shaders for mip level generation',
          code: `
            struct VSOutput {
              @builtin(position) position: vec4f,
              @location(0) texcoord: vec2f,
            };

            @vertex fn vs(
              @builtin(vertex_index) vertexIndex : u32
            ) -> VSOutput {
              let pos = array(

                vec2f( 0.0,  0.0),  // center
                vec2f( 1.0,  0.0),  // right, center
                vec2f( 0.0,  1.0),  // center, top

                // 2st triangle
                vec2f( 0.0,  1.0),  // center, top
                vec2f( 1.0,  0.0),  // right, center
                vec2f( 1.0,  1.0),  // right, top
              );

              var vsOutput: VSOutput;
              let xy = pos[vertexIndex];
              vsOutput.position = vec4f(xy * 2.0 - 1.0, 0.0, 1.0);
              vsOutput.texcoord = vec2f(xy.x, 1.0 - xy.y);
              return vsOutput;
            }

            @group(0) @binding(0) var ourSampler: sampler;
            @group(0) @binding(1) var ourTexture: texture_2d<f32>;

            @fragment fn fs(fsInput: VSOutput) -> @location(0) vec4f {
              return textureSample(ourTexture, ourSampler, fsInput.texcoord);
            }
          `,
        });

        sampler = device.createSampler({
          minFilter: 'linear',
        });
      }

      const encoder = device.createCommandEncoder({
        label: 'mip gen encoder',
      });

        const format = texture.format;
        if (!pipelineByFormat[format]) {
          pipelineByFormat[format] = device.createRenderPipeline({
            label: 'mip level generator pipeline',
            layout: 'auto',
            vertex: {
              module,
            },
            fragment: {
              module,
              targets: [{ format }],
            },
          });
        }
        const pipeline = pipelineByFormat[format];

        for (let baseMipLevel = 1; baseMipLevel < texture.mipLevelCount; ++baseMipLevel) {
          const bindGroup = device.createBindGroup({
            layout: pipeline.getBindGroupLayout(0),
            entries: [
              { binding: 0, resource: sampler },
              { binding: 1, resource: texture.createView({baseMipLevel: baseMipLevel - 1, mipLevelCount: 1}) },
            ],
          });

          const renderPassDescriptor = {
            label: 'our basic canvas renderPass',
            colorAttachments: [
              {
                view: texture.createView({baseMipLevel, mipLevelCount: 1}),
                loadOp: 'clear',
                storeOp: 'store',
              },
            ],
          };

          const pass = encoder.beginRenderPass(renderPassDescriptor);
          pass.setPipeline(pipeline);
          pass.setBindGroup(0, bindGroup);
          pass.draw(6);  // call our vertex shader 6 times
          pass.end();
        }

      const commandBuffer = encoder.finish();
      device.queue.submit([commandBuffer]);
    };
  })();
*/

  const generateMips = (() => {
    let sampler;
    const resourcesByFormat = {};

    return function generateMips(device, texture) {
      const resources = resourcesByFormat[texture.format] ?? (() => {

        sampler = sampler ?? device.createSampler({
          minFilter: 'linear',
          magFilter: 'linear',
        });

        const dummyTextureViews = range(3, () => device.createTexture({
          format: 'rgba8unorm',
          size: [1],
          usage: GPUTextureUsage.STORAGE_BINDING
        }).createView());

        const module = device.createShaderModule({
          code: `
            @group(0) @binding(0) var smp: sampler;
            @group(0) @binding(1) var mip0: texture_2d<f32>;
            @group(0) @binding(2) var mip1: texture_storage_2d<rgba8unorm, write>;
            @group(0) @binding(3) var mip2: texture_storage_2d<rgba8unorm, write>;
            @group(0) @binding(4) var mip3: texture_storage_2d<rgba8unorm, write>;
            @group(0) @binding(5) var mip4: texture_storage_2d<rgba8unorm, write>;

            var<workgroup> texels1: array<array<vec4f, 8>, 8>;
            var<workgroup> texels2: array<array<vec4f, 4>, 4>;
            var<workgroup> texels3: array<array<vec4f, 2>, 2>;

            // It doesn't seem like we need to check bounds. We bind a dummy texture.
            // for each mip level not used. We use textureSampleLevel for the top level
            // so it will clamp-to-edge. textureStore is speced to "not execute" if
            // out of bounds.

            fn processMip0ToMip1(blockXY: vec2u, lid: vec2u) {
              let mip1TexelXY = blockXY * 8 + lid.xy;
              let texelNdx = lid.xy;
              let mip1Size = textureDimensions(mip1);

              let uv = (vec2f(mip1TexelXY) + 0.5) / vec2f(mip1Size);
              let c = textureSampleLevel(mip0, smp, uv, 0.0);
              texels1[texelNdx.y][texelNdx.x] = c;
              textureStore(mip1, mip1TexelXY, c);
            }

            @compute @workgroup_size(8, 8) fn cs(
              @builtin(local_invocation_id) lid: vec3u,
              @builtin(workgroup_id) wid: vec3u,
            ) {
              let blockXY = wid.xy;

              // generate mip1 from mip0
              processMip0ToMip1(blockXY, lid.xy);
              workgroupBarrier();

              // generate mip2 from mip1
              if (lid.x < 4 && lid.y < 4) {
                let mip2TexelXY = blockXY * 4 + lid.xy;
                let texelNdx = lid.xy;
                let srcNdx = texelNdx * 2;
                let c0 = texels1[srcNdx.y    ][srcNdx.x    ];
                let c1 = texels1[srcNdx.y    ][srcNdx.x + 1];
                let c2 = texels1[srcNdx.y + 1][srcNdx.x    ];
                let c3 = texels1[srcNdx.y + 1][srcNdx.x + 1];
                let c = mix(mix(c0, c1, 0.5), mix(c2, c3, 0.5), 0.5);
                texels2[texelNdx.y][texelNdx.x] = c;
                textureStore(mip2, mip2TexelXY, c);
              }

              workgroupBarrier();

              // generate mip3 from mip2
              if (lid.x < 2 && lid.y < 2) {
                let mip3TexelXY = blockXY * 2 + lid.xy;
                let texelNdx = lid.xy;
                let srcNdx = texelNdx * 2;
                let c0 = texels2[srcNdx.y    ][srcNdx.x    ];
                let c1 = texels2[srcNdx.y    ][srcNdx.x + 1];
                let c2 = texels2[srcNdx.y + 1][srcNdx.x    ];
                let c3 = texels2[srcNdx.y + 1][srcNdx.x + 1];
                let c = mix(mix(c0, c1, 0.5), mix(c2, c3, 0.5), 0.5);
                texels3[texelNdx.y][texelNdx.x] = c;
                textureStore(mip3, mip3TexelXY, c);
              }

              workgroupBarrier();

              // generate mip4 from mip3
              if (lid.x < 1 && lid.y < 1) {
                let mip4TexelXY = blockXY + lid.xy;
                let texelNdx = lid.xy;
                let srcNdx = texelNdx * 2;
                let c0 = texels3[srcNdx.y    ][srcNdx.x    ];
                let c1 = texels3[srcNdx.y    ][srcNdx.x + 1];
                let c2 = texels3[srcNdx.y + 1][srcNdx.x    ];
                let c3 = texels3[srcNdx.y + 1][srcNdx.x + 1];
                let c = mix(mix(c0, c1, 0.5), mix(c2, c3, 0.5), 0.5);
                textureStore(mip4, mip4TexelXY, c);
              }

            }
          `,
        });

        const pipeline = device.createComputePipeline({
          layout: 'auto',
          compute: { module },
        });

        return {
          pipeline,
          dummyTextureViews,
        };
      })();
      resourcesByFormat[texture.format] = resources;
      const { pipeline, dummyTextureViews } = resources;

      const encoder = device.createCommandEncoder();
      const pass = encoder.beginComputePass();
      pass.setPipeline(pipeline);

      // 0: 4096  ---0
      // 1: 2048  ---1
      // 2: 1024  ---2
      // 3:  512  ---3
      // 4:  256  ---4  ---0
      // 5:  128        ---1
      // 6:   64        ---2
      // 7:   32        ---3
      // 8:   16        ---4 ---0
      // 9:    8             ---1
      //10:    4             ---2
      //11:    2             ---3
      //12:    1             ---4

      const getMipLevelView = (texture, baseMipLevel, dummyNdx) => baseMipLevel < texture.mipLevelCount
        ? texture.createView({ baseMipLevel: baseMipLevel, mipLevelCount: 1 })
        : dummyTextureViews[dummyNdx];

      const kLevelsPerPass = 4;
      const kBlockSize = 16;
      const numMipLevelsToGenerate = texture.mipLevelCount - 1;
      for (let baseMipLevel = 0; baseMipLevel < numMipLevelsToGenerate; baseMipLevel += kLevelsPerPass) {
        const levelSize = [
          Math.max(1, texture.width >> baseMipLevel),
          Math.max(1, texture.height >> baseMipLevel),
        ];
        const bindGroup = device.createBindGroup({
          layout: pipeline.getBindGroupLayout(0),
          entries: [
            { binding: 0, resource: sampler },
            { binding: 1, resource: texture.createView({ baseMipLevel: baseMipLevel, mipLevelCount: 1 }) },
            { binding: 2, resource: texture.createView({ baseMipLevel: baseMipLevel + 1, mipLevelCount: 1 }) },
            { binding: 3, resource: getMipLevelView(texture, baseMipLevel + 2, 0) },
            { binding: 4, resource: getMipLevelView(texture, baseMipLevel + 3, 1) },
            { binding: 5, resource: getMipLevelView(texture, baseMipLevel + 4, 2) },
          ],
        });

        pass.setBindGroup(0, bindGroup);
        pass.dispatchWorkgroups(...levelSize.map(size => Math.ceil(size / kBlockSize)));
      }
      pass.end();
      device.queue.submit([encoder.finish()]);      
    }    
  })();

  const size = 2048;
  const half = size / 2;

  const texture = device.createTexture({
    size: [size, size],
    format: 'rgba8unorm',
    usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.TEXTURE_BINDING,
  });
  const view = texture.createView();

  // offsets to the various uniform values in float32 indices
  const kMatrixOffset = 0;

  const sampler = device.createSampler();

  // create a buffer for the uniform values
  const uniformBufferSize =
    16 * 4; // matrix is 16 32bit floats (4bytes each)
  const uniformBuffer = device.createBuffer({
    label: 'uniforms for quad',
    size: uniformBufferSize,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });

  // create a typedarray to hold the values for the uniforms in JavaScript
  const uniformValues = new Float32Array(uniformBufferSize / 4);
  const matrix = uniformValues.subarray(kMatrixOffset, 16);

  const bindGroup = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: sampler },
      { binding: 1, resource: texture.createView() },
      { binding: 2, resource: { buffer: uniformBuffer }},
    ],
  });

  const renderPassDescriptor = {
    label: 'our basic canvas renderPass',
    colorAttachments: [
      {
        // view: <- to be filled out when we render
        clearValue: [0.3, 0.3, 0.3, 1],
        loadOp: 'clear',
        storeOp: 'store',
      },
    ],
  };

  let then = 0;
  const info = document.querySelector('#info');
  const settings = {
    count: 1,
  };
  const gui = new GUI();
  gui.add(settings, 'count', 1, 1000, 1);

  function render(time) {
    const deltaTime = time - then;
    then = time;
    fpsAverage.addSample(deltaTime);
    info.textContent = `fps: ${(1000 / fpsAverage.get()).toFixed(0)}`;

    // Get the current texture from the canvas context and
    // set it as the texture to render to.
    const canvasTexture = context.getCurrentTexture();
    renderPassDescriptor.colorAttachments[0].view = canvasTexture.createView();

    //renderPassDescriptor.colorAttachments[0].view = texture.createView();
    renderPassDescriptor.colorAttachments[0].clearValue[0] = time % 1;

    for (let i = 0; i < settings.count; ++i) {
      {
        const encoder = device.createCommandEncoder({});
        renderPassDescriptor.colorAttachments[0].view = view;
        const pass = encoder.beginRenderPass(renderPassDescriptor);
        pass.end();
        const commandBuffer = encoder.finish();
        device.queue.submit([commandBuffer]);
      }

      generateMips(device, texture);

      {
        const encoder = device.createCommandEncoder({});
        renderPassDescriptor.colorAttachments[0].view = canvasTexture.createView();
        const pass = encoder.beginRenderPass(renderPassDescriptor);
        pass.setPipeline(pipeline);
        pass.setBindGroup(0, bindGroup);
        pass.draw(6);  // call our vertex shader 6 times
        pass.end();

        const commandBuffer = encoder.finish();
        device.queue.submit([commandBuffer]);
      }
    }

    requestAnimationFrame(render);
  }
  requestAnimationFrame(render);

  const observer = new ResizeObserver(entries => {
    for (const entry of entries) {
      const canvas = entry.target;
      const width = entry.contentBoxSize[0].inlineSize;
      const height = entry.contentBoxSize[0].blockSize;
      canvas.width = Math.max(1, Math.min(width, device.limits.maxTextureDimension2D));
      canvas.height = Math.max(1, Math.min(height, device.limits.maxTextureDimension2D));
    }
  });
  observer.observe(canvas);


  canvas.addEventListener('click', () => {
    texNdx = (texNdx + 1) % textures.length;
  });
}

function fail(msg) {
  // eslint-disable-next-line no-alert
  alert(msg);
}

main();
  

jsGist.json

{"name":"WebGPU Texture - Generate Mipmaps Speed Test (compute mipmap gen)","settings":{},"filenames":["index.html","index.css","index.js"]}