Method init_shapes

src/opr/impl/basic_arith.cpp:1084–1188

}

void Reduce::KernScheduler::init_shapes(
        megdnn::Reduce* opr, CompNode comp_node, DType inp_dtype, Mode mode,
        TensorShape ishp, TensorShape oshp, const Param::DataType data_type) {
    mgb_assert(ishp.ndim && oshp.ndim);

    if (ishp.eq_shape(m_prev_ishp) && oshp.eq_shape(m_prev_oshp))
        return;

    m_prev_ishp = ishp;
    m_prev_oshp = oshp;

    m_kern_param.clear();

    if (oshp.is_scalar()) {
        // if ishp is non-contiguous, add_layout_constraint_contiguous would be
        // added; so we do not have to worry about this
        ishp.shape[0] = ishp.total_nr_elems();
        ishp.ndim = 1;
    }

    mgb_assert(
            oshp.ndim == ishp.ndim,
            "input and output ndim mismatch for reduction: ishp=%s oshp=%s",
            ishp.to_string().c_str(), oshp.to_string().c_str());

    for (size_t i = 0; i < ishp.ndim; ++i) {
        if (ishp.shape[i] != oshp.shape[i]) {
            mgb_assert(
                    oshp.shape[i] == 1,
                    "input and output shape mismatch for reduction: "
                    "ishp=%s oshp=%s",
                    ishp.to_string().c_str(), oshp.to_string().c_str());
        }
    }

    auto remove_axis = [](TensorShape& shp, size_t ax) {
        mgb_assert(shp.ndim > 1);
        for (auto i = ax + 1; i < shp.ndim; ++i)
            shp.shape[i - 1] = shp.shape[i];
        --shp.ndim;
    };

    // collapse consecutive shape-1 axes in oshp
    for (size_t i = 0; i < oshp.ndim; ++i) {
        auto start = i;
        while (i < oshp.ndim && oshp.shape[i] == 1)
            ++i;

        if (start + 1 < i) {
            for (auto j = start + 1; j < i; ++j)
                ishp.shape[start] *= ishp.shape[j];

            for (auto j = start + 1; j < i; ++j) {
                remove_axis(ishp, start + 1);
                remove_axis(oshp, start + 1);
            }

            i = start;