div.comp 776 B

123456789101112131415161718192021222324252627
  1. #version 450
  2. #include "types.comp"
  3. #include "generic_binary_head.comp"
  4. const uint num_threads = 256;
  5. layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in;
  6. void main() {
  7. uint idx = get_idx();
  8. // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation
  9. const uint num_iter = 2;
  10. [[unroll]] for (uint i = 0; i < num_iter; ++i) {
  11. if (idx >= p.ne) {
  12. continue;
  13. }
  14. uint i00, i01, i02, i03;
  15. get_indices(idx, i00, i01, i02, i03);
  16. data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) / FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)]));
  17. idx += num_threads;
  18. }
  19. }