| 656 | |
| 657 | template<typename Ti, typename Tw, typename To> |
| 658 | To mean_all(Param<Ti> in) { |
| 659 | using std::unique_ptr; |
| 660 | uintl in_elements = |
| 661 | in.info.dims[0] * in.info.dims[1] * in.info.dims[2] * in.info.dims[3]; |
| 662 | bool is_linear = (in.info.strides[0] == 1); |
| 663 | for (int k = 1; k < 4; k++) { |
| 664 | is_linear &= (in.info.strides[k] == |
| 665 | (in.info.strides[k - 1] * in.info.dims[k - 1])); |
| 666 | } |
| 667 | |
| 668 | // FIXME: Use better heuristics to get to the optimum number |
| 669 | if (in_elements > 4096 || !is_linear) { |
| 670 | if (is_linear) { |
| 671 | in.info.dims[0] = in_elements; |
| 672 | for (int k = 1; k < 4; k++) { |
| 673 | in.info.dims[k] = 1; |
| 674 | in.info.strides[k] = in_elements; |
| 675 | } |
| 676 | } |
| 677 | |
| 678 | uint threads_x = nextpow2(std::max(32u, (uint)in.info.dims[0])); |
| 679 | threads_x = std::min(threads_x, THREADS_PER_BLOCK); |
| 680 | uint threads_y = THREADS_PER_BLOCK / threads_x; |
| 681 | |
| 682 | uint blocks_x = divup(in.info.dims[0], threads_x * REPEAT); |
| 683 | uint blocks_y = divup(in.info.dims[1], threads_y); |
| 684 | |
| 685 | dim4 outDims(blocks_x, in.info.dims[1], in.info.dims[2], |
| 686 | in.info.dims[3]); |
| 687 | |
| 688 | Array<To> tmpOut = createEmptyArray<To>(outDims); |
| 689 | Array<Tw> tmpCt = createEmptyArray<Tw>(outDims); |
| 690 | |
| 691 | Param<Tw> iwt; |
| 692 | mean_first_launcher<Ti, Tw, To>(tmpOut, tmpCt, in, iwt, blocks_x, |
| 693 | blocks_y, threads_x); |
| 694 | |
| 695 | uintl tmp_elements = tmpOut.elements(); |
| 696 | |
| 697 | compute_t<To> val; |
| 698 | auto tmpOut_get = tmpOut.get(); |
| 699 | auto tmpCt_get = tmpCt.get(); |
| 700 | getQueue() |
| 701 | .submit([&](sycl::handler &h) { |
| 702 | auto out = |
| 703 | tmpOut_get->get_host_access(h, sycl::read_only); |
| 704 | auto ct = |
| 705 | tmpCt_get->get_host_access(h, sycl::read_only); |
| 706 | |
| 707 | h.host_task([out, ct, tmp_elements, &val] { |
| 708 | val = static_cast<compute_t<To>>(out[0]); |
| 709 | compute_t<Tw> weight = static_cast<compute_t<Tw>>(ct[0]); |
| 710 | |
| 711 | for (int i = 1; i < tmp_elements; i++) { |
| 712 | stable_mean(&val, &weight, compute_t<To>(out[i]), |
| 713 | compute_t<Tw>(ct[i])); |
| 714 | } |
| 715 | }); |