WebGPU: Test WGSL convertCubeCoordToNormalized3DTextureCoord / convertNormalized3DTexCoordToCubeCoord

README.md

WebGPU: Test WGSL convertCubeCoordToNormalized3DTextureCoord / convertNormalized3DTexCoordToCubeCoord

view on jsgist

index.css

:root {
  color-scheme: light dark;
}
pre { margin: 0; }

index.html

/*bug-in-github-api-content-can-not-be-empty*/

index.js

import {
  getSizeAndAlignmentOfUnsizedArrayElement,
  makeShaderDataDefinitions,
  makeStructuredView,
} from 'https://greggman.github.io/webgpu-utils/dist/1.x/webgpu-utils.module.js';

const assert = (cond, msg = '') => {
  if (!cond) { throw new Error(msg); }
};

const code = `
fn convertCubeCoordToNormalized3DTextureCoord(v: vec3f) -> vec3f {
  var uvw: vec3f;
  var layer: f32;
  let r = normalize(v);
  let absR = abs(r);
  if (absR.x > absR.y && absR.x > absR.z) {
    // x major
    if (r.x >= 0.0) {
      uvw = vec3f(-r.z, -r.y, absR.x);
      layer = 0;
    } else {
      uvw = vec3f(r.z, -r.y, absR.x);
      layer = 1;
    }
  } else if (absR.y > absR.z) {
    // y major
    if (r.y >= 0.0) {
      uvw = vec3f(r.x, r.z, absR.y);
      layer = 2;
    } else {
      uvw = vec3f(r.x, -r.z, absR.y);
      layer = 3;
    }
  } else {
    // z major
    if (r.z >= 0.0) {
      uvw = vec3f(r.x, -r.y, absR.z);
      layer = 4;
    } else {
      uvw = vec3f(-r.x, -r.y, absR.z);
      layer = 5;
    }
  }
  return vec3f((uvw.xy / uvw.z + 1.0) * 0.5, (layer + 0.5) / 6.0);
}

const faceMat = array(
  mat3x3f( 0,  0,  -2,  0, -2,   0,  1,  1,   1),   // pos-x
  mat3x3f( 0,  0,   2,  0, -2,   0, -1,  1,  -1),   // neg-x
  mat3x3f( 2,  0,   0,  0,  0,   2, -1,  1,  -1),   // pos-y
  mat3x3f( 2,  0,   0,  0,  0,  -2, -1, -1,   1),   // neg-y
  mat3x3f( 2,  0,   0,  0, -2,   0, -1,  1,   1),   // pos-z
  mat3x3f(-2,  0,   0,  0, -2,   0,  1,  1,  -1));  // neg-z

fn convertNormalized3DTexCoordToCubeCoord(uvLayer: vec3f) -> vec3f {
  let layer = u32(uvLayer.z * 6.0);
  return normalize(faceMat[layer] * vec3f(uvLayer.xy, 1));
}

struct Info {
  cc0: vec3f,
  uvl0: vec3f,
  cc1: vec3f,
  uvl1: vec3f,
  cc2: vec3f,
  matches: u32,
}

@group(0) @binding(0) var<storage, read_write> data: array<Info>;

@compute @workgroup_size(1) fn cs(
    @builtin(global_invocation_id) ndx: vec3u,
    @builtin(num_workgroups) size: vec3u) {
  let cc0 = (vec3f(ndx) / 2.0 + 0.25) / vec3f(size * 2) * 2.0 - 1.0;
  let uvl0 = convertCubeCoordToNormalized3DTextureCoord(cc0);
  let cc1 = convertNormalized3DTexCoordToCubeCoord(uvl0);
  let uvl1 = convertCubeCoordToNormalized3DTextureCoord(cc1);
  let cc2 = convertNormalized3DTexCoordToCubeCoord(uvl1);

  let index = ndx.z * size.x * size.y + ndx.y * size.x + ndx.x;
  data[index] = Info(cc0, uvl0, cc1, uvl1, cc2, 0);
}

@group(0) @binding(0) var<storage, read_write> texels: array<vec4f>;
@group(0) @binding(1) var s: sampler;
@group(0) @binding(2) var t: texture_cube<f32>;

@compute @workgroup_size(1) fn readTexture(
    @builtin(global_invocation_id) ndx: vec3u,
    @builtin(num_workgroups) size: vec3u) {
  let index = ndx.z * size.x * size.y + ndx.y * size.x + ndx.x;
  let tSize = textureDimensions(t, 0);
  let uvl = vec3f((vec2f(ndx.xy) + 0.5) / vec2f(tSize), 11.0 / 12.0);
  let cc = convertNormalized3DTexCoordToCubeCoord(uvl);
  let uvl1 = convertCubeCoordToNormalized3DTextureCoord(cc);
  let cc2 = convertNormalized3DTexCoordToCubeCoord(uvl1);
  texels[index] = textureSampleLevel(t, s, cc2, 0) * 255.0;
}
`;

const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();
device.addEventListener('uncapturederror', e => {
  e.preventDefault(); console.error(e.error.message);
});

const texture = device.createTexture({
  size: [8, 8, 6],
  mipLevelCount: 3,
  usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.TEXTURE_BINDING,
  textureBindingViewDimension: 'cube',
  format: 'rgba8unorm',
});
const texels = new Uint8Array(texture.width * texture.height * 4);
for (let f = 0; f < texture.depthOrArrayLayers; ++f) {
  for (let y = 0; y < texture.height; ++y) {
    for (let x = 0; x < texture.width; ++x) {
      const off = (y * texture.width + x) * 4;
      texels[off + 0] = x;
      texels[off + 1] = y;
      texels[off + 2] = f;
      texels[off + 3] = (y + x) % 2 ? 0 : 255;
    }
  }
  device.queue.writeTexture(
    { texture, origin: [0, 0, f] },
    texels,
    { bytesPerRow: texture.width * 4 },
    [ texture.width, texture.height ],
  );
}

const sampler = device.createSampler({
  minFilter: 'linear',
  magFilter: 'linear',
  mipmapFilter: 'linear',
})

if (true) {
  const module = device.createShaderModule({ code });
  const pipeline = device.createComputePipeline({
    layout: 'auto',
    compute: { module, entryPoint: 'readTexture' },
  });
  const storageBuffer = device.createBuffer({
    size: texture.width * texture.height * 4 * 4,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
  });
  const resultBuffer = device.createBuffer({
    size: storageBuffer.size,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
  });

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

  const encoder = device.createCommandEncoder();
  const pass = encoder.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(texture.width, texture.height);
  pass.end();
  encoder.copyBufferToBuffer(storageBuffer, 0, resultBuffer, 0, resultBuffer.size);
  device.queue.submit([encoder.finish()]);

  await resultBuffer.mapAsync(GPUMapMode.READ);
  const results = new Float32Array(resultBuffer.getMappedRange());

  for (let y = 0; y < texture.height; ++y) {
    log('--- y:', y);
    for (let x = 0; x < texture.width; ++x) {
      const offset = (y * texture.width + x) * 4;
      log(`x: ${x}: ${results[offset + 0]}, ${results[offset + 1]}, ${results[offset + 2]}, ${results[offset + 3]}`)
    }
  }
}

if (false) {
  const module = device.createShaderModule({ code });
  const pipeline = device.createComputePipeline({
    layout: 'auto',
    compute: { module, entryPoint: 'cs' },
  });
  const defs = makeShaderDataDefinitions(code);
  const {size} = getSizeAndAlignmentOfUnsizedArrayElement(defs.storages.data);
  const kDims = 4;
  const storageBuffer = device.createBuffer({
    size: size * kDims * kDims * kDims,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
  });
  const resultBuffer = device.createBuffer({
    size: storageBuffer.size,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
  });

  const bindGroup = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: storageBuffer }},
    ],
  });

  const encoder = device.createCommandEncoder();
  const pass = encoder.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(kDims, kDims, kDims);
  pass.end();
  encoder.copyBufferToBuffer(storageBuffer, 0, resultBuffer, 0, resultBuffer.size);
  device.queue.submit([encoder.finish()]);

  await resultBuffer.mapAsync(GPUMapMode.READ);
  const { views } = makeStructuredView(defs.storages.data, resultBuffer.getMappedRange());


  const kFaceUVMatrices/*: mat3[]*/ =
    /* prettier-ignore */ [
    [ 0,  0,  -2,  0, -2,   0,  1,  1,   1],   // pos-x
    [ 0,  0,   2,  0, -2,   0, -1,  1,  -1],   // neg-x
    [ 2,  0,   0,  0,  0,   2, -1,  1,  -1],   // pos-y
    [ 2,  0,   0,  0,  0,  -2, -1, -1,   1],   // neg-y
    [ 2,  0,   0,  0, -2,   0, -1,  1,   1],   // pos-z
    [-2,  0,   0,  0, -2,   0,  1,  1,  -1],   // neg-z
  ];

  /** multiply a vec3 by mat3 */
  function transformMat3(v/*: vec3*/, m/*: mat3*/)/*: vec3*/ {
    const x = v[0];
    const y = v[1];
    const z = v[2];

    return [
      x * m[0] + y * m[3] + z * m[6],
      x * m[1] + y * m[4] + z * m[7],
      x * m[2] + y * m[5] + z * m[8],
    ];
  }

  /** normalize a vec3 */
  function normalize(v/*: vec3*/)/*: vec3*/ {
    const length = Math.sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
    assert(length > 0);
    return v.map(v => v / length)/* as vec3*/;
  }

  /**
   * Converts a cube map coordinate to a uv coordinate (0 to 1) and layer (0.5/6.0 to 5.5/6.0).
   */
  function convertCubeCoordToNormalized3DTextureCoord(v/*: vec3*/)/*: vec3*/ {
    let uvw;
    let layer;
    // normalize the coord.
    // MAINTENANCE_TODO: handle(0, 0, 0)
    const r = normalize(v);
    const absR = r.map(v => Math.abs(v));
    if (absR[0] > absR[1] && absR[0] > absR[2]) {
      // x major
      const negX = r[0] < 0.0 ? 1 : 0;
      uvw = [negX ? r[2] : -r[2], -r[1], absR[0]];
      layer = negX;
    } else if (absR[1] > absR[2]) {
      // y major
      const negY = r[1] < 0.0 ? 1 : 0;
      uvw = [r[0], negY ? -r[2] : r[2], absR[1]];
      layer = 2 + negY;
    } else {
      // z major
      const negZ = r[2] < 0.0 ? 1 : 0;
      uvw = [negZ ? -r[0] : r[0], -r[1], absR[2]];
      layer = 4 + negZ;
    }
    return [(uvw[0] / uvw[2] + 1) * 0.5, (uvw[1] / uvw[2] + 1) * 0.5, (layer + 0.5) / 6];
  }

  /**
   * Convert a 3d texcoord into a cube map coordinate.
   */
  function convertNormalized3DTexCoordToCubeCoord(uvLayer/*: vec3*/) {
    const [u, v, faceLayer] = uvLayer;
    return normalize(transformMat3([u, v, 1], kFaceUVMatrices[Math.min(5, faceLayer * 6) | 0]));
  }


  const p3 = v => v.toFixed(3);
  const v3 = v => `${p3(v[0])},${p3(v[1])},${p3(v[2])}`

  for (let z = 0; z < kDims; ++z) {
    log('---z:', z);
    for (let y = 0; y < kDims; ++y) {
      log('  ---y:', y);
      for (let x = 0; x < kDims; ++x) {
        const index = z * kDims * kDims + y * kDims + x;
        const info = views[index];

        /*
        let cc0 = (vec3f(ndx) / 2.0 + 0.25) / vec3f(size * 2) * 2.0 - 1.0;
        let uvl0 = convertCubeCoordToNormalized3DTextureCoord(cc0);
        let cc1 = convertNormalized3DTexCoordToCubeCoord(uvl0);
        let uvl1 = convertCubeCoordToNormalized3DTextureCoord(cc1);
        let cc2 = convertNormalized3DTexCoordToCubeCoord(uvl1);
        */
        const cc0 = [x, y, z].map(v => (v / 2 + 0.25) / (kDims * 2) * 2 - 1);
        const uvl0 = convertCubeCoordToNormalized3DTextureCoord(cc0);
        const cc1 = convertNormalized3DTexCoordToCubeCoord(uvl0);
        const uvl1 = convertCubeCoordToNormalized3DTextureCoord(cc1);
        const cc2 = convertNormalized3DTexCoordToCubeCoord(uvl1);

        log(`cpu   ${x},${y},${z}: cc0: ${v3(cc0)}, uvl0: ${v3(uvl0)} cc1: ${v3(cc1)} uvl1: ${v3(uvl1)} cc2: ${v3(cc2)}`);
        log(`gpu   ${x},${y},${z}: cc0: ${v3(info.cc0)}, uvl0: ${v3(info.uvl0)} cc1: ${v3(info.cc1)} uvl1: ${v3(info.uvl1)} cc2: ${v3(info.cc2)}`);
      }
    }
  }
}

device.queue.submit([]);

function log(...args) {
  const elem = document.createElement('pre');
  elem.textContent = args.join(' ');
  document.body.appendChild(elem);
}

jsGist.json

{"name":"WebGPU: Test WGSL convertCubeCoordToNormalized3DTextureCoord / convertNormalized3DTexCoordToCubeCoord","settings":{},"filenames":["index.html","index.css","index.js"]}