| 562 | } // anonymous namespace |
| 563 | |
| 564 | void ReduceImpl::exec( |
| 565 | _megdnn_tensor_in src, _megdnn_tensor_out dst, _megdnn_workspace workspace) { |
| 566 | check_exec(src.layout, dst.layout, workspace.size); |
| 567 | size_t A, B, C; |
| 568 | reduce::get_ABC(src.layout, A, B, C, param().axis); |
| 569 | bool execed = false; |
| 570 | using Mode = param::Reduce::Mode; |
| 571 | #define DISPATCH_FUNC(Reducer, _dtype, ctype, comp_type) \ |
| 572 | if (C == 1) { \ |
| 573 | using _Reducer = Reducer<_dtype, ctype, comp_type, true>; \ |
| 574 | using _ReducerC1SmallB = Reducer<_dtype, ctype, comp_type, false>; \ |
| 575 | std::function<void(const ctype*, ctype*, DType, size_t, size_t, size_t)> \ |
| 576 | do_reduce = Exec<_Reducer, true>::do_reduce; \ |
| 577 | if (src.layout.dtype.category() != DTypeCategory::FLOAT) { \ |
| 578 | if (B == 2) \ |
| 579 | do_reduce = ExecC1SmallB<_ReducerC1SmallB, ctype, 2>::do_reduce; \ |
| 580 | if (B == 3) \ |
| 581 | do_reduce = ExecC1SmallB<_ReducerC1SmallB, ctype, 3>::do_reduce; \ |
| 582 | if (B == 4) \ |
| 583 | do_reduce = ExecC1SmallB<_ReducerC1SmallB, ctype, 4>::do_reduce; \ |
| 584 | } \ |
| 585 | MIDOUT_BEGIN( \ |
| 586 | megdnn_arm_common_reduce, ctype, _dtype, comp_type, midout_iv(1)) { \ |
| 587 | MEGDNN_DISPATCH_CPU_KERN_OPR(do_reduce( \ |
| 588 | reinterpret_cast<ctype*>(src.raw_ptr()), \ |
| 589 | reinterpret_cast<ctype*>(dst.raw_ptr()), src_type, A, B, C)); \ |
| 590 | execed = true; \ |
| 591 | } \ |
| 592 | MIDOUT_END(); \ |
| 593 | } else { \ |
| 594 | using _Reducer = Reducer<_dtype, ctype, comp_type, false>; \ |
| 595 | std::function<void(const ctype*, ctype*, DType, size_t, size_t, size_t)> \ |
| 596 | do_reduce = Exec<_Reducer, false>::do_reduce; \ |
| 597 | MIDOUT_BEGIN( \ |
| 598 | megdnn_arm_common_reduce, ctype, _dtype, comp_type, midout_iv(1)) { \ |
| 599 | MEGDNN_DISPATCH_CPU_KERN_OPR(do_reduce( \ |
| 600 | reinterpret_cast<ctype*>(src.raw_ptr()), \ |
| 601 | reinterpret_cast<ctype*>(dst.raw_ptr()), src_type, A, B, C)); \ |
| 602 | execed = true; \ |
| 603 | } \ |
| 604 | MIDOUT_END(); \ |
| 605 | } |
| 606 | |
| 607 | #define DISPATCH_MODE_QUANTIZED(dtype, ctype, comp_type) \ |
| 608 | switch (param().mode) { \ |
| 609 | case Mode::MEAN: \ |
| 610 | DISPATCH_FUNC(MeanReducer, dtype, ctype, comp_type); \ |
| 611 | break; \ |
| 612 | case Mode::MAX: \ |
| 613 | DISPATCH_FUNC(maxReducer, dtype, ctype, ctype); \ |
| 614 | break; \ |
| 615 | case Mode::MIN: \ |
| 616 | DISPATCH_FUNC(minReducer, dtype, ctype, ctype); \ |
| 617 | break; \ |
| 618 | default: \ |
| 619 | break; \ |
| 620 | } |
| 621 | |