dawn-cmake/test/bug/tint/943.spvasm.expected.wgsl

[[block]]
struct Uniforms {
  NAN : f32;
  [[size(12)]]
  padding : u32;
  aShape : vec3<i32>;
  [[size(4)]]
  padding_1 : u32;
  bShape : vec3<i32>;
  [[size(4)]]
  padding_2 : u32;
  outShape : vec3<i32>;
  [[size(4)]]
  padding_3 : u32;
  outShapeStrides : vec2<i32>;
};

type RTArr = [[stride(4)]] array<f32>;

type RTArr_1 = [[stride(4)]] array<f32>;

[[block]]
struct ssbOut {
  result : RTArr_1;
};

type RTArr_2 = [[stride(4)]] array<f32>;

[[block]]
struct ssbA {
  A : RTArr_1;
};

[[block]]
struct ssbB {
  B : RTArr_1;
};

var<private> dimAOuter_1 : i32;

[[group(0), binding(3)]] var<uniform> x_48 : Uniforms;

var<private> dimInner_1 : i32;

var<private> dimBOuter_1 : i32;

[[group(0), binding(0)]] var<storage, read_write> x_54 : ssbOut;

var<private> gl_LocalInvocationID : vec3<u32>;

var<private> gl_GlobalInvocationID : vec3<u32>;

var<workgroup> mm_Asub : array<array<f32, 64u>, 64u>;

var<workgroup> mm_Bsub : array<array<f32, 1u>, 64u>;

[[group(0), binding(1)]] var<storage, read> x_165 : ssbA;

var<private> batch : i32;

[[group(0), binding(2)]] var<storage, read> x_185 : ssbB;

fn coordsInBounds_vi2_vi2_(coord : ptr<function, vec2<i32>>, shape : ptr<function, vec2<i32>>) -> bool {
  var x_87 : bool;
  var x_88_phi : bool;
  let x_76 : vec2<i32> = *(coord);
  let x_81 : bool = all((x_76 >= vec2<i32>(0, 0)));
  x_88_phi = x_81;
  if (x_81) {
    let x_84 : vec2<i32> = *(coord);
    let x_85 : vec2<i32> = *(shape);
    x_87 = all((x_84 < x_85));
    x_88_phi = x_87;
  }
  let x_88 : bool = x_88_phi;
  return x_88;
}

fn mm_readA_i1_i1_(row : ptr<function, i32>, col : ptr<function, i32>) -> f32 {
  var batchASize : i32;
  var param_10 : vec2<i32>;
  var param_11 : vec2<i32>;
  var x_430 : f32;
  let x_417 : i32 = x_48.aShape.y;
  let x_419 : i32 = x_48.aShape.z;
  batchASize = (x_417 * x_419);
  let x_421 : i32 = *(row);
  let x_422 : i32 = *(col);
  let x_424 : i32 = dimAOuter_1;
  let x_425 : i32 = dimInner_1;
  param_10 = vec2<i32>(x_421, x_422);
  param_11 = vec2<i32>(x_424, x_425);
  let x_429 : bool = coordsInBounds_vi2_vi2_(&(param_10), &(param_11));
  if (x_429) {
    let x_438 : i32 = batch;
    let x_439 : i32 = batchASize;
    let x_441 : i32 = *(row);
    let x_442 : i32 = dimInner_1;
    let x_445 : i32 = *(col);
    let x_448 : f32 = x_165.A[(((x_438 * x_439) + (x_441 * x_442)) + x_445)];
    x_430 = x_448;
  } else {
    x_430 = 0.0;
  }
  let x_450 : f32 = x_430;
  return x_450;
}

fn mm_readB_i1_i1_(row_1 : ptr<function, i32>, col_1 : ptr<function, i32>) -> f32 {
  var batchBSize : i32;
  var param_12 : vec2<i32>;
  var param_13 : vec2<i32>;
  var x_468 : f32;
  let x_455 : i32 = x_48.bShape.y;
  let x_457 : i32 = x_48.bShape.z;
  batchBSize = (x_455 * x_457);
  let x_459 : i32 = *(row_1);
  let x_460 : i32 = *(col_1);
  let x_462 : i32 = dimInner_1;
  let x_463 : i32 = dimBOuter_1;
  param_12 = vec2<i32>(x_459, x_460);
  param_13 = vec2<i32>(x_462, x_463);
  let x_467 : bool = coordsInBounds_vi2_vi2_(&(param_12), &(param_13));
  if (x_467) {
    let x_475 : i32 = batch;
    let x_476 : i32 = batchBSize;
    let x_478 : i32 = *(row_1);
    let x_479 : i32 = dimBOuter_1;
    let x_482 : i32 = *(col_1);
    let x_485 : f32 = x_185.B[(((x_475 * x_476) + (x_478 * x_479)) + x_482)];
    x_468 = x_485;
  } else {
    x_468 = 0.0;
  }
  let x_487 : f32 = x_468;
  return x_487;
}

fn getOutputFlatIndex_vi3_(coords : ptr<function, vec3<i32>>) -> i32 {
  let x_99 : vec3<i32> = *(coords);
  let x_105 : i32 = x_48.outShapeStrides.x;
  let x_107 : i32 = x_48.outShapeStrides.y;
  return i32(dot(vec3<f32>(x_99), vec3<f32>(vec3<i32>(x_105, x_107, 1))));
}

fn setOutput_i1_f1_(flatIndex : ptr<function, i32>, value : ptr<function, f32>) {
  let x_95 : i32 = *(flatIndex);
  let x_96 : f32 = *(value);
  x_54.result[x_95] = x_96;
  return;
}

fn setOutput_i1_i1_i1_f1_(d0 : ptr<function, i32>, d1 : ptr<function, i32>, d2 : ptr<function, i32>, value_1 : ptr<function, f32>) {
  var flatIndex_1 : i32;
  var param : vec3<i32>;
  var param_1 : i32;
  var param_2 : f32;
  let x_115 : i32 = *(d0);
  let x_116 : i32 = *(d1);
  let x_117 : i32 = *(d2);
  param = vec3<i32>(x_115, x_116, x_117);
  let x_120 : i32 = getOutputFlatIndex_vi3_(&(param));
  flatIndex_1 = x_120;
  let x_122 : i32 = flatIndex_1;
  param_1 = x_122;
  let x_124 : f32 = *(value_1);
  param_2 = x_124;
  setOutput_i1_f1_(&(param_1), &(param_2));
  return;
}

fn mm_write_i1_i1_f1_(row_2 : ptr<function, i32>, col_2 : ptr<function, i32>, value_2 : ptr<function, f32>) {
  var outCoord : vec3<i32>;
  var param_14 : i32;
  var param_15 : i32;
  var param_16 : i32;
  var param_17 : f32;
  let x_491 : i32 = batch;
  let x_492 : i32 = *(row_2);
  let x_493 : i32 = *(col_2);
  outCoord = vec3<i32>(x_491, x_492, x_493);
  let x_496 : i32 = batch;
  param_14 = x_496;
  let x_498 : i32 = *(row_2);
  param_15 = x_498;
  let x_500 : i32 = *(col_2);
  param_16 = x_500;
  let x_502 : f32 = *(value_2);
  param_17 = x_502;
  setOutput_i1_i1_i1_f1_(&(param_14), &(param_15), &(param_16), &(param_17));
  return;
}

fn mm_matMul_i1_i1_i1_(dimAOuter : ptr<function, i32>, dimInner : ptr<function, i32>, dimBOuter : ptr<function, i32>) {
  var tileRow : i32;
  var tileCol : i32;
  var globalRow : i32;
  var globalCol : i32;
  var numTiles : i32;
  var innerRow : i32;
  var innerCol : i32;
  var acc : array<array<f32, 1u>, 1u>;
  var tileColA : i32;
  var tileRowB : i32;
  var t : i32;
  var innerRow_1 : i32;
  var innerCol_1 : i32;
  var inputRow : i32;
  var inputCol : i32;
  var param_3 : i32;
  var param_4 : i32;
  var innerRow_2 : i32;
  var innerCol_2 : i32;
  var inputRow_1 : i32;
  var inputCol_1 : i32;
  var param_5 : i32;
  var param_6 : i32;
  var k : i32;
  var inner : i32;
  var BCached : array<f32, 1u>;
  var innerRow_3 : i32;
  var ACached : f32;
  var innerCol_3 : i32;
  var innerRow_4 : i32;
  var innerCol_4 : i32;
  var param_7 : i32;
  var param_8 : i32;
  var param_9 : f32;
  let x_132 : u32 = gl_LocalInvocationID.y;
  tileRow = (bitcast<i32>(x_132) * 1);
  let x_137 : u32 = gl_LocalInvocationID.x;
  tileCol = (bitcast<i32>(x_137) * 1);
  let x_143 : u32 = gl_GlobalInvocationID.y;
  globalRow = (bitcast<i32>(x_143) * 1);
  let x_148 : u32 = gl_GlobalInvocationID.x;
  globalCol = (bitcast<i32>(x_148) * 1);
  let x_152 : i32 = *(dimInner);
  numTiles = (((x_152 - 1) / 64) + 1);
  innerRow = 0;
  loop {
    let x_163 : i32 = innerRow;
    if ((x_163 < 1)) {
    } else {
      break;
    }
    innerCol = 0;
    loop {
      let x_171 : i32 = innerCol;
      if ((x_171 < 1)) {
      } else {
        break;
      }
      let x_177 : i32 = innerRow;
      let x_178 : i32 = innerCol;
      acc[x_177][x_178] = 0.0;

      continuing {
        let x_181 : i32 = innerCol;
        innerCol = (x_181 + 1);
      }
    }

    continuing {
      let x_183 : i32 = innerRow;
      innerRow = (x_183 + 1);
    }
  }
  let x_187 : u32 = gl_LocalInvocationID.x;
  tileColA = (bitcast<i32>(x_187) * 64);
  let x_192 : u32 = gl_LocalInvocationID.y;
  tileRowB = (bitcast<i32>(x_192) * 1);
  t = 0;
  loop {
    let x_201 : i32 = t;
    let x_202 : i32 = numTiles;
    if ((x_201 < x_202)) {
    } else {
      break;
    }
    innerRow_1 = 0;
    loop {
      let x_210 : i32 = innerRow_1;
      if ((x_210 < 1)) {
      } else {
        break;
      }
      innerCol_1 = 0;
      loop {
        let x_218 : i32 = innerCol_1;
        if ((x_218 < 64)) {
        } else {
          break;
        }
        let x_221 : i32 = tileRow;
        let x_222 : i32 = innerRow_1;
        inputRow = (x_221 + x_222);
        let x_225 : i32 = tileColA;
        let x_226 : i32 = innerCol_1;
        inputCol = (x_225 + x_226);
        let x_233 : i32 = inputRow;
        let x_234 : i32 = inputCol;
        let x_235 : i32 = globalRow;
        let x_236 : i32 = innerRow_1;
        let x_238 : i32 = t;
        let x_240 : i32 = inputCol;
        param_3 = (x_235 + x_236);
        param_4 = ((x_238 * 64) + x_240);
        let x_244 : f32 = mm_readA_i1_i1_(&(param_3), &(param_4));
        mm_Asub[x_233][x_234] = x_244;

        continuing {
          let x_247 : i32 = innerCol_1;
          innerCol_1 = (x_247 + 1);
        }
      }

      continuing {
        let x_249 : i32 = innerRow_1;
        innerRow_1 = (x_249 + 1);
      }
    }
    innerRow_2 = 0;
    loop {
      let x_257 : i32 = innerRow_2;
      if ((x_257 < 1)) {
      } else {
        break;
      }
      innerCol_2 = 0;
      loop {
        let x_265 : i32 = innerCol_2;
        if ((x_265 < 1)) {
        } else {
          break;
        }
        let x_268 : i32 = tileRowB;
        let x_269 : i32 = innerRow_2;
        inputRow_1 = (x_268 + x_269);
        let x_272 : i32 = tileCol;
        let x_273 : i32 = innerCol_2;
        inputCol_1 = (x_272 + x_273);
        let x_278 : i32 = inputRow_1;
        let x_279 : i32 = inputCol_1;
        let x_280 : i32 = t;
        let x_282 : i32 = inputRow_1;
        let x_284 : i32 = globalCol;
        let x_285 : i32 = innerCol_2;
        param_5 = ((x_280 * 64) + x_282);
        param_6 = (x_284 + x_285);
        let x_289 : f32 = mm_readB_i1_i1_(&(param_5), &(param_6));
        mm_Bsub[x_278][x_279] = x_289;

        continuing {
          let x_291 : i32 = innerCol_2;
          innerCol_2 = (x_291 + 1);
        }
      }

      continuing {
        let x_293 : i32 = innerRow_2;
        innerRow_2 = (x_293 + 1);
      }
    }
    workgroupBarrier();
    k = 0;
    loop {
      let x_302 : i32 = k;
      if ((x_302 < 64)) {
      } else {
        break;
      }
      inner = 0;
      loop {
        let x_310 : i32 = inner;
        if ((x_310 < 1)) {
        } else {
          break;
        }
        let x_314 : i32 = inner;
        let x_315 : i32 = k;
        let x_316 : i32 = tileCol;
        let x_317 : i32 = inner;
        let x_320 : f32 = mm_Bsub[x_315][(x_316 + x_317)];
        BCached[x_314] = x_320;

        continuing {
          let x_322 : i32 = inner;
          inner = (x_322 + 1);
        }
      }
      innerRow_3 = 0;
      loop {
        let x_330 : i32 = innerRow_3;
        if ((x_330 < 1)) {
        } else {
          break;
        }
        let x_333 : i32 = tileRow;
        let x_334 : i32 = innerRow_3;
        let x_336 : i32 = k;
        let x_338 : f32 = mm_Asub[(x_333 + x_334)][x_336];
        ACached = x_338;
        innerCol_3 = 0;
        loop {
          let x_345 : i32 = innerCol_3;
          if ((x_345 < 1)) {
          } else {
            break;
          }
          let x_347 : i32 = innerRow_3;
          let x_348 : i32 = innerCol_3;
          let x_349 : f32 = ACached;
          let x_350 : i32 = innerCol_3;
          let x_352 : f32 = BCached[x_350];
          let x_355 : f32 = acc[x_347][x_348];
          acc[x_347][x_348] = (x_355 + (x_349 * x_352));

          continuing {
            let x_358 : i32 = innerCol_3;
            innerCol_3 = (x_358 + 1);
          }
        }

        continuing {
          let x_360 : i32 = innerRow_3;
          innerRow_3 = (x_360 + 1);
        }
      }

      continuing {
        let x_362 : i32 = k;
        k = (x_362 + 1);
      }
    }
    workgroupBarrier();

    continuing {
      let x_364 : i32 = t;
      t = (x_364 + 1);
    }
  }
  innerRow_4 = 0;
  loop {
    let x_372 : i32 = innerRow_4;
    if ((x_372 < 1)) {
    } else {
      break;
    }
    innerCol_4 = 0;
    loop {
      var x_393 : bool;
      var x_394_phi : bool;
      let x_380 : i32 = innerCol_4;
      if ((x_380 < 1)) {
      } else {
        break;
      }
      let x_382 : i32 = globalCol;
      let x_383 : i32 = innerCol_4;
      let x_385 : i32 = *(dimBOuter);
      let x_386 : bool = ((x_382 + x_383) < x_385);
      x_394_phi = x_386;
      if (x_386) {
        let x_389 : i32 = globalRow;
        let x_390 : i32 = innerRow_4;
        let x_392 : i32 = *(dimAOuter);
        x_393 = ((x_389 + x_390) < x_392);
        x_394_phi = x_393;
      }
      let x_394 : bool = x_394_phi;
      if (x_394) {
        let x_397 : i32 = globalRow;
        let x_398 : i32 = innerRow_4;
        let x_400 : i32 = globalCol;
        let x_401 : i32 = innerCol_4;
        let x_403 : i32 = innerRow_4;
        let x_404 : i32 = innerCol_4;
        param_7 = (x_397 + x_398);
        param_8 = (x_400 + x_401);
        let x_409 : f32 = acc[x_403][x_404];
        param_9 = x_409;
        mm_write_i1_i1_f1_(&(param_7), &(param_8), &(param_9));
      }

      continuing {
        let x_411 : i32 = innerCol_4;
        innerCol_4 = (x_411 + 1);
      }
    }

    continuing {
      let x_413 : i32 = innerRow_4;
      innerRow_4 = (x_413 + 1);
    }
  }
  return;
}

fn main_1() {
  var param_18 : i32;
  var param_19 : i32;
  var param_20 : i32;
  let x_67 : i32 = x_48.aShape.y;
  dimAOuter_1 = x_67;
  let x_71 : i32 = x_48.aShape.z;
  dimInner_1 = x_71;
  let x_75 : i32 = x_48.bShape.z;
  dimBOuter_1 = x_75;
  let x_505 : u32 = gl_GlobalInvocationID.z;
  batch = bitcast<i32>(x_505);
  let x_508 : i32 = dimAOuter_1;
  param_18 = x_508;
  let x_510 : i32 = dimInner_1;
  param_19 = x_510;
  let x_512 : i32 = dimBOuter_1;
  param_20 = x_512;
  mm_matMul_i1_i1_i1_(&(param_18), &(param_19), &(param_20));
  return;
}

[[stage(compute), workgroup_size(1, 64, 1)]]
fn main([[builtin(local_invocation_id)]] gl_LocalInvocationID_param : vec3<u32>, [[builtin(global_invocation_id)]] gl_GlobalInvocationID_param : vec3<u32>) {
  gl_LocalInvocationID = gl_LocalInvocationID_param;
  gl_GlobalInvocationID = gl_GlobalInvocationID_param;
  main_1();
}