Function computeWorkPerThreadForConv2d

tfjs-backend-webgpu/src/webgpu_util.ts:120–140

(
    layout: {x: number[], y?: number[], z?: number[]}, outputShape: number[],
    isVec4 = false)

}

export function computeWorkPerThreadForConv2d(
    layout: {x: number[], y?: number[], z?: number[]}, outputShape: number[],
    isVec4 = false): [number, number, number] {
  if (isVec4) {
    return [4, 4, 1];
  }

  const dim0 = arrayProduct(layout.x.map(d => outputShape[d]));
  const dim1 = arrayProduct(layout.y.map(d => outputShape[d]));
  // TODO(jiajia.qin@intel.com): More fine tune based on outputShape.
  // The following conditions correspond to the values set in
  // computeWorkgroupSizeForConv2d.
  if (dim0 <= 4) {
    return [1, 2, 1];
  }
  if (dim1 <= 4) {
    return [2, 1, 1];
  }

  return [2, 2, 1];
}

export function flatDispatchLayout(shape: number[]) {
  return {x: shape.map((d, i) => i)};