import { c as a } from "../../webgpu_util-DuofJBMo.js";
import { e as l } from "../../webgpu_program-CbjdYLYk.js";
import { v as p } from "../../index-DSGwv2Yx.js";
class n {
  variableNames = ["A"];
  outputShape;
  shaderKey;
  dispatchLayout;
  dispatch;
  // Note that the maximum number of workgroup invocations by webgpu is 256.
  // Nick: Reduce to 8x8
  workgroupSize = [8, 8, 1];
  constructor(t, o) {
    const i = t.length, e = new Array(i), r = t.slice();
    r[r.length - 1] *= 2;
    for (let u = 0; u < e.length; u++)
      e[u] = r[o[u]];
    e[e.length - 1] /= 2, this.outputShape = e, this.dispatchLayout = i === 2 ? { x: [0], y: [1] } : { x: [1], y: [2], z: [0] }, this.dispatch = a(this.dispatchLayout, this.outputShape, this.workgroupSize, [2, 1, 1]), this.shaderKey = `transposeShared16_${i}`;
  }
  getUserCode() {
    const t = this.outputShape.length;
    p(
      this.workgroupSize[0] === this.workgroupSize[1],
      () => `Must be a square tile, current tile shape is ${this.workgroupSize[0]} x ${this.workgroupSize[1]}`
    );
    const o = this.workgroupSize[0] * 2;
    return `
      var<workgroup> tile : array<array<f32, ${o + 1}>, ${o}>;
      ${l()} {
        var x = i32(workgroupId.x) * ${o / 2} + i32(localId.x);
        var y = i32(workgroupId.y) * ${o} + i32(localId.y);
        let batch = ${t === 3 ? "i32(workgroupId.z)" : "0"};
        let batchOffsetA = ${t === 3 ? "batch * uniforms.aShapeStrides[0]" : "0"};
        let batchOffsetOut = ${t === 3 ? "batch * uniforms.outShapeStrides[0]" : "0"};

        let inputWidth = uniforms.outShape[${t === 3 ? "1" : "0"}] / 2; // Output height
        let inputHeight = uniforms.outShape[${t === 3 ? "2" : "1"}] * 2; // Output width
        if (x < inputWidth && y < inputHeight) {
            let unpackedA = unpack2x16float(u32(A[batchOffsetA + y * inputWidth + x]));
            tile[localId.y][localId.x * 2] = unpackedA.x;
            tile[localId.y][localId.x * 2 + 1] = unpackedA.y;
        }
        // Second load to cover the tile
        y = y + ${this.workgroupSize[0]};
        if (x < inputWidth && y < inputHeight) {
            let unpackedA = unpack2x16float(u32(A[batchOffsetA + y * inputWidth + x]));
            tile[localId.y + ${this.workgroupSize[0]}][localId.x * 2] = unpackedA.x;
            tile[localId.y + ${this.workgroupSize[0]}][localId.x * 2 + 1] = unpackedA.y;
        }
        workgroupBarrier();

        let outputWidth = uniforms.outShape[${t === 3 ? "2" : "1"}]; // Output width
        let outputHeight = uniforms.outShape[${t === 3 ? "1" : "0"}] * 2; // Output height
        x = i32(workgroupId.y) * ${o / 2} + i32(localId.x);
        y = i32(workgroupId.x) * ${o} + i32(localId.y);
        if (x < outputWidth && y < outputHeight) {
          result[batchOffsetOut + y * outputWidth + x] = i32(pack2x16float(vec2<f32>(tile[localId.x * 2][localId.y], tile[localId.x * 2 + 1][localId.y])));
        }
        // Second store to cover the tile
        y = y + ${this.workgroupSize[0]};
        if (x < outputWidth && y < outputHeight) {
          result[batchOffsetOut + y * outputWidth + x] = i32(pack2x16float(vec2<f32>(tile[localId.x * 2][localId.y + ${this.workgroupSize[0]}], tile[localId.x * 2 + 1][localId.y + ${this.workgroupSize[0]}])));
        }
      }
    `;
  }
}
export {
  n as default
};
