| 565 | |
| 566 | template<typename T, typename Tw> |
| 567 | T mean_all_weighted(Param<T> in, Param<Tw> iwt) { |
| 568 | uintl in_elements = |
| 569 | in.info.dims[0] * in.info.dims[1] * in.info.dims[2] * in.info.dims[3]; |
| 570 | // FIXME: Use better heuristics to get to the optimum number |
| 571 | if (in_elements > 4096) { |
| 572 | bool in_is_linear = (in.info.strides[0] == 1); |
| 573 | bool wt_is_linear = (iwt.info.strides[0] == 1); |
| 574 | for (int k = 1; k < 4; k++) { |
| 575 | in_is_linear &= (in.info.strides[k] == |
| 576 | (in.info.strides[k - 1] * in.info.dims[k - 1])); |
| 577 | wt_is_linear &= (iwt.info.strides[k] == |
| 578 | (iwt.info.strides[k - 1] * iwt.info.dims[k - 1])); |
| 579 | } |
| 580 | |
| 581 | if (in_is_linear && wt_is_linear) { |
| 582 | in.info.dims[0] = in_elements; |
| 583 | for (int k = 1; k < 4; k++) { |
| 584 | in.info.dims[k] = 1; |
| 585 | in.info.strides[k] = in_elements; |
| 586 | } |
| 587 | |
| 588 | for (int k = 0; k < 4; k++) { |
| 589 | iwt.info.dims[k] = in.info.dims[k]; |
| 590 | iwt.info.strides[k] = in.info.strides[k]; |
| 591 | } |
| 592 | } |
| 593 | |
| 594 | uint threads_x = nextpow2(std::max(32u, (uint)in.info.dims[0])); |
| 595 | threads_x = std::min(threads_x, THREADS_PER_BLOCK); |
| 596 | uint threads_y = THREADS_PER_BLOCK / threads_x; |
| 597 | |
| 598 | uint blocks_x = divup(in.info.dims[0], threads_x * REPEAT); |
| 599 | uint blocks_y = divup(in.info.dims[1], threads_y); |
| 600 | |
| 601 | Array<T> tmpOut = createEmptyArray<T>( |
| 602 | {blocks_x, in.info.dims[1], in.info.dims[2], in.info.dims[3]}); |
| 603 | Array<Tw> tmpWt = createEmptyArray<Tw>( |
| 604 | {blocks_x, in.info.dims[1], in.info.dims[2], in.info.dims[3]}); |
| 605 | |
| 606 | uintl tmp_elements = tmpOut.elements(); |
| 607 | |
| 608 | mean_first_launcher<T, Tw, T>(tmpOut, tmpWt, in, iwt, blocks_x, |
| 609 | blocks_y, threads_x); |
| 610 | |
| 611 | compute_t<T> val; |
| 612 | auto tmpOut_get = tmpOut.get(); |
| 613 | auto tmpWt_get = tmpWt.get(); |
| 614 | getQueue() |
| 615 | .submit([&](sycl::handler &h) { |
| 616 | auto acc_in = |
| 617 | tmpOut_get->get_host_access(h, sycl::read_only); |
| 618 | auto acc_wt = |
| 619 | tmpWt_get->get_host_access(h, sycl::read_only); |
| 620 | |
| 621 | h.host_task([acc_in, acc_wt, tmp_elements, &val] { |
| 622 | val = static_cast<compute_t<T>>(acc_in[0]); |
| 623 | compute_t<Tw> weight = |
| 624 | static_cast<compute_t<Tw>>(acc_wt[0]); |