struct MatrixBatch {
  size: vec4<u32>,
  data: array<f32>,
}

struct FVectorBatch {
  size: vec4<u32>,
  data: array<f32>,
}

struct UVectorBatch {
  size: vec4<u32>,
  data: array<u32>,
}

@group(0) @binding(0)
var<storage, read> m: MatrixBatch;
@group(0) @binding(1)
var<storage, read> r: FVectorBatch;
@group(0) @binding(2)
var<storage, read_write> result: UVectorBatch;

@compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  result.size.x = m.size.y;
  result.size.z = m.size.z;
  var resi = global_id.z * result.size.x + global_id.x;
  var tp = r.data[global_id.z * r.size.x + global_id.x];
  var p = 0f;
  var mi = (global_id.z * m.size.y + global_id.x) * m.size.x;
  for (var i = 0u; i < m.size.x; i++) {
    p += m.data[mi];
    if (p >= tp) {
      result.data[resi] = i;
      break;
    }
    mi++;
  }
}