Function reduce_raw

src/opr/test/basic_arith/reduction.cpp:98–153

template <Mode mode, typename ctype>
void reduce_raw(HostTensorND& dest, const HostTensorND& src) {
    auto tshp = dest.shape();
    using Impl = ImplTrait<mode, ctype>;

    if (tshp.is_scalar()) {
        if (src.shape().is_scalar()) {
            dest.copy_from_fixlayout(src);
            return;
        }

        Impl impl;
        ctype val = impl.init();
        for (auto i : megdnn::tensor_iter_valonly<ctype>(src.as_megdnn()))
            val = impl.reduce(val, i);
        dest.ptr<ctype>()[0] = impl.finalize(val);
        return;
    }

    mgb_assert(tshp.ndim == src.shape().ndim);

    std::vector<size_t> axis_to_use;
    for (size_t i = 0; i < tshp.ndim; i++) {
        if (tshp.shape[i] != src.shape(i)) {
            mgb_assert(tshp.shape[i] == 1);
            axis_to_use.push_back(i);
        }
    }

    if (axis_to_use.empty()) {
        dest.copy_from_fixlayout(src);
        return;
    }
    TensorLayout sub_layout{dest.dtype()};
    sub_layout.ndim = axis_to_use.size();
    for (size_t i = 0; i < axis_to_use.size(); i++) {
        sub_layout.shape[i] = src.layout().shape[axis_to_use[i]];
        sub_layout.stride[i] = src.layout().stride[axis_to_use[i]];
    }

    auto diter_maker = megdnn::tensor_iter<ctype>(dest.as_megdnn());
    for (auto iter = diter_maker.begin(), iter_end = diter_maker.end();
         iter != iter_end; ++iter) {
        ptrdiff_t offset = 0;
        for (size_t i = 0; i < tshp.ndim; i++)
            offset += iter.idx()[i] * src.layout().stride[i];

        Impl impl;
        ctype val = impl.init();
        auto subspec = SubTensorSpec::make_from_offset_elem(sub_layout, offset);
        HostTensorND subt = const_cast<HostTensorND&>(src).sub(subspec);
        for (ctype i : megdnn::tensor_iter_valonly<ctype>(subt.as_megdnn())) {
            val = impl.reduce(val, i);
        }
        *iter = impl.finalize(val);
    }
}

template <Mode mode, class dtype>