in src/webgpu/shader/execution/expression/call/builtin/texture_utils.ts [369:571]
export async function queryMipLevelMixWeightsForDevice(t: GPUTest, stage: ShaderStage) {
const { device } = t;
const kNumWeightTypes = 2;
assert(kNumWeightTypes <= 4);
const module = device.createShaderModule({
code: `
@group(0) @binding(0) var tex: texture_2d<f32>;
@group(0) @binding(1) var smp: sampler;
@group(0) @binding(2) var<storage, read_write> result: array<vec4f>;
struct VSOutput {
@builtin(position) pos: vec4f,
@location(0) @interpolate(flat, either) ndx: u32,
@location(1) @interpolate(flat, either) result: vec4f,
};
fn getMixLevels(wNdx: u32) -> vec4f {
let mipLevel = f32(wNdx) / ${kMipLevelWeightSteps};
let size = textureDimensions(tex);
let g = mix(1.0, 2.0, mipLevel) / f32(size.x);
let ddx = vec2f(g, 0);
return vec4f(
textureSampleLevel(tex, smp, vec2f(0.5), mipLevel).r,
textureSampleGrad(tex, smp, vec2f(0.5), ddx, vec2f(0)).r,
0,
0);
}
fn getPosition(vNdx: u32) -> vec4f {
let pos = array(
vec2f(-1, 3),
vec2f( 3, -1),
vec2f(-1, -1),
);
let p = pos[vNdx];
return vec4f(p, 0, 1);
}
// -- for getting fragment stage weights --
@vertex fn vs(@builtin(vertex_index) vNdx: u32, @builtin(instance_index) iNdx: u32) -> VSOutput {
return VSOutput(getPosition(vNdx), iNdx, vec4f(0));
}
@fragment fn fsRecord(v: VSOutput) -> @location(0) vec4u {
return bitcast<vec4u>(getMixLevels(v.ndx));
}
// -- for getting compute stage weights --
@compute @workgroup_size(1) fn csRecord(@builtin(global_invocation_id) id: vec3u) {
result[id.x] = getMixLevels(id.x);
}
// -- for getting vertex stage weights --
@vertex fn vsRecord(@builtin(vertex_index) vNdx: u32, @builtin(instance_index) iNdx: u32) -> VSOutput {
return VSOutput(getPosition(vNdx), iNdx, getMixLevels(iNdx));
}
@fragment fn fsSaveVs(v: VSOutput) -> @location(0) vec4u {
return bitcast<vec4u>(v.result);
}
`,
});
const texture = t.createTextureTracked({
size: [2, 2, 1],
format: 'r8unorm',
usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
mipLevelCount: 2,
});
device.queue.writeTexture(
{ texture, mipLevel: 1 },
new Uint8Array([255]),
{ bytesPerRow: 1 },
[1, 1]
);
const sampler = device.createSampler({
minFilter: 'linear',
magFilter: 'linear',
mipmapFilter: 'linear',
});
const target = t.createTextureTracked({
size: [kMipLevelWeightSteps + 1, 1],
format: 'rgba32uint',
usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_SRC,
});
const storageBuffer = t.createBufferTracked({
label: 'queryMipLevelMixWeightsForDevice:storageBuffer',
size: 4 * 4 * (kMipLevelWeightSteps + 1),
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
});
const resultBuffer = t.createBufferTracked({
label: 'queryMipLevelMixWeightsForDevice:resultBuffer',
size: align(storageBuffer.size, 256), // padded for copyTextureToBuffer
usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
});
const createBindGroup = (pipeline: GPUComputePipeline | GPURenderPipeline) =>
device.createBindGroup({
layout: pipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: texture.createView() },
{ binding: 1, resource: sampler },
...(stage === 'compute' ? [{ binding: 2, resource: { buffer: storageBuffer } }] : []),
],
});
const encoder = device.createCommandEncoder({ label: 'queryMipLevelMixWeightsForDevice' });
switch (stage) {
case 'compute': {
const pipeline = device.createComputePipeline({
layout: 'auto',
compute: { module },
});
const pass = encoder.beginComputePass();
pass.setPipeline(pipeline);
pass.setBindGroup(0, createBindGroup(pipeline));
pass.dispatchWorkgroups(kMipLevelWeightSteps + 1);
pass.end();
encoder.copyBufferToBuffer(storageBuffer, 0, resultBuffer, 0, storageBuffer.size);
break;
}
case 'fragment': {
const pipeline = device.createRenderPipeline({
layout: 'auto',
vertex: { module, entryPoint: 'vs' },
fragment: { module, entryPoint: 'fsRecord', targets: [{ format: 'rgba32uint' }] },
});
const pass = encoder.beginRenderPass({
colorAttachments: [
{
view: target.createView(),
loadOp: 'clear',
storeOp: 'store',
},
],
});
pass.setPipeline(pipeline);
pass.setBindGroup(0, createBindGroup(pipeline));
for (let x = 0; x <= kMipLevelWeightSteps; ++x) {
pass.setViewport(x, 0, 1, 1, 0, 1);
pass.draw(3, 1, 0, x);
}
pass.end();
encoder.copyTextureToBuffer({ texture: target }, { buffer: resultBuffer }, [target.width]);
break;
}
case 'vertex': {
const pipeline = device.createRenderPipeline({
layout: 'auto',
vertex: { module, entryPoint: 'vsRecord' },
fragment: { module, entryPoint: 'fsSaveVs', targets: [{ format: 'rgba32uint' }] },
});
const pass = encoder.beginRenderPass({
colorAttachments: [
{
view: target.createView(),
loadOp: 'clear',
storeOp: 'store',
},
],
});
pass.setPipeline(pipeline);
pass.setBindGroup(0, createBindGroup(pipeline));
for (let x = 0; x <= kMipLevelWeightSteps; ++x) {
pass.setViewport(x, 0, 1, 1, 0, 1);
pass.draw(3, 1, 0, x);
}
pass.end();
encoder.copyTextureToBuffer({ texture: target }, { buffer: resultBuffer }, [target.width]);
break;
}
}
device.queue.submit([encoder.finish()]);
await resultBuffer.mapAsync(GPUMapMode.READ);
// need to map a sub-portion since we may have padded the buffer.
const result = Array.from(
new Float32Array(resultBuffer.getMappedRange(0, (kMipLevelWeightSteps + 1) * 16))
);
resultBuffer.unmap();
resultBuffer.destroy();
const [sampleLevelWeights, gradWeights] = unzip(result, kNumWeightTypes, 4);
validateWeights(t, stage, sampleLevelWeights);
validateWeights(t, stage, gradWeights);
texture.destroy();
storageBuffer.destroy();
return {
sampleLevelWeights,
softwareMixToGPUMixGradWeights: generateSoftwareMixToGPUMixGradWeights(gradWeights, texture),
};
}