Function symbolicallyComputeStrides

tfjs-backend-webgpu/src/shader_util.ts:19–35

(
    indicesArr: number[], variableName: string)

// Generates WGSL that computes strides.
export function symbolicallyComputeStrides(
    indicesArr: number[], variableName: string): string[] {
  if (Math.max(...indicesArr) > 5) {
    throw new Error('Cannot symbolically compute strides for rank > 6 tensor.');
  }

  const numCoords = indicesArr.length;
  const indicesStr = 'xyzwuv';
  const shape = indicesArr.map(d => `${variableName}.${indicesStr[d]}`);
  const strides = new Array(numCoords - 1);
  strides[numCoords - 2] = shape[numCoords - 1];
  for (let i = numCoords - 3; i >= 0; --i) {
    strides[i] = `(${strides[i + 1]} * ${shape[i + 1]})`;
  }

  return strides;
}

export const atomicAddSnippet =
    (ptr: string, v: string, type: 'int32'|'float32') => {