onednn-src 0.1.13

/*******************************************************************************
* Copyright 2022 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#include "gpu/intel/conv/jit/normalization.hpp"
#include "gpu/intel/conv/jit/config.hpp"
#include "gpu/intel/jit/ir/tensor.hpp"

namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace conv {
namespace jit {

layout_t insert_dimension(const layout_t &layout, const pvar_t &idx) {
    auto new_blocks = layout.blocks();
    for (auto &b : new_blocks) {
        if (b.idx.index() >= idx.index()) b.idx = b.idx + 1;
    }
    return layout_t(layout.type(), new_blocks, layout.offset(),
            layout.ndims() + 1,
            /*do_normalize=*/false);
}

layout_t remove_size_1_dimension(const layout_t &layout, const pvar_t &idx) {
    gpu_assert(idx.index() < layout.ndims());
    gpu_assert(layout.elems(idx) == 1);
    dim_assignment_t a(layout.ndims(), layout.ndims() - 1);
    for (dim_idx_t i = 0; i < layout.ndims(); i++) {
        if (i == idx.index()) continue;
        a.assign(i, i < idx.index() ? i : i - 1);
    }
    return a.map(layout);
}

layout_t split_dimension(
        const layout_t &_layout, const pvar_t &idx, dim_t outer_size) {
    dim_t rem_inner_size
            = ir_utils::safe_divide(_layout.elems(idx), outer_size);
    auto layout = insert_dimension(_layout, idx);
    std::vector<layout_block_t> new_blocks;
    for (auto &b : layout.blocks()) {
        if (b.idx.index() != idx.index() + 1) {
            new_blocks.push_back(b);
            continue;
        }
        if (b.size % rem_inner_size == 0) {
            new_blocks.emplace_back(idx.index() + 1, rem_inner_size, b.stride);
            new_blocks.emplace_back(idx, b.size / rem_inner_size,
                    dim_t(b.stride) * rem_inner_size);
            rem_inner_size = 1;
        } else {
            new_blocks.push_back(b);
            rem_inner_size = ir_utils::safe_divide(rem_inner_size, b.size);
        }
    }

    // Remove inner blocks with size one.
    std::vector<layout_block_t> _new_blocks;
    pvar_map_t<bool> seen;
    for (auto it = new_blocks.rbegin(); it != new_blocks.rend(); ++it) {
        if (it->size == 1 && seen[it->idx]) continue;
        _new_blocks.push_back(*it);
        seen[it->idx] = true;
    }
    std::reverse(_new_blocks.begin(), _new_blocks.end());
    return layout.with(_new_blocks, false);
}

layout_t normalize_conv_groups(const layout_t &layout, bool with_groups,
        dim_t groups, bool is_dw, bool add_groups, bool is_wei) {
    if (with_groups == add_groups) return layout;
    if (is_wei) {
        gpu_assert(groups == 1)
                << "Adding/removing groups can be done only for single group.";
        if (add_groups) return insert_dimension(layout, 0);
        return remove_size_1_dimension(layout, 0);
    }

    gpu_assert(!with_groups) << "Unexpected groups in source/destination.";
    if (is_dw) groups = layout.elems(1);
    if (layout.elems(1) == 1) groups = 1;
    return split_dimension(layout, /*dim=*/1, groups);
}

layout_t normalize_layout(const layout_t &_layout, bool with_groups,
        dim_t groups, bool is_dw, const std::array<int, 3> &dhw_map,
        bool add_groups, bool is_wei) {
    layout_t layout = _layout;
    layout = spatials_to_3d(layout, with_groups, dhw_map);
    layout = normalize_conv_groups(
            layout, with_groups, groups, is_dw, add_groups, is_wei);

    return layout;
}

tile_t normalize_tile(std::vector<dim_t> &dims, bool with_groups, dim_t groups,
        bool is_dw, const std::array<int, 3> &dhw_map, bool add_groups,
        bool is_wei) {
    layout_t dummy_layout(dsl::type_t::u8(), dims);
    return normalize_layout(dummy_layout, with_groups, groups, is_dw, dhw_map,
            add_groups, is_wei)
            .tile();
}

void normalize_layouts(layout_t &src_layout, layout_t &wei_layout,
        layout_t &dst_layout, layout_t &bia_layout, bool with_groups, dim_t g,
        dim_t ic, dim_t oc, bool is_dw, const std::array<int, 3> &dhw_map,
        bool add_groups) {
    src_layout = normalize_layout(src_layout, /*with_groups=*/false,
            g > 1 ? src_layout.elems(1) / ic : 1, is_dw, dhw_map, add_groups,
            /*is_wei=*/false);
    wei_layout = normalize_layout(wei_layout, with_groups, g, is_dw, dhw_map,
            add_groups, /*is_wei=*/true);
    dst_layout = normalize_layout(dst_layout, /*with_groups=*/false,
            g > 1 ? dst_layout.elems(1) / oc : 1, is_dw, dhw_map, add_groups,
            /*is_wei=*/false);
    if (add_groups && !bia_layout.is_empty()) {
        gpu_assert(bia_layout.ndims() == 1) << bia_layout;
        bia_layout = split_dimension(bia_layout, 0, g);
    }
}

uint32_t post_op_view_mapper_t::normalize_mask(uint32_t orig_mask) const {
    dim_idx_t cp_ndims = cp_view().nvdims();
    gpu_assert(cp_ndims >= 3);
    // Add groups to match ngcdhw layout.
    bool add_groups = (cp_view().vvars()[1].as<var_t>().name == "g");
    // Number of dimensions before normalization.
    int orig_ndims = prb_.ndims;
    std::vector<dim_t> dummy_dims(orig_ndims, 1);
    dim_t mask_set_value = 2;
    for (int i = 0; i < orig_ndims; i++) {
        if ((orig_mask & (1 << i)) != 0) dummy_dims[i] = mask_set_value;
    }
    auto cvt_dims = normalize_tile(dummy_dims, /*with_groups=*/false, prb_.g,
            prb_.is_dw, prb_.dhw_map,
            /*add_groups=*/false, /*is_wei=*/false);
    // Split channels into groups and channels to match ngcdhw layout.
    if (add_groups) {
        tile_t new_tile;
        for (auto &b : cvt_dims) {
            if (b > 0) new_tile.set(b + 1, cvt_dims[b]);
        }
        new_tile[1] = cvt_dims[1];
        cvt_dims = std::move(new_tile);
    }

    uint32_t mask = 0;
    for (dim_idx_t i = 0; i < cp_ndims; i++) {
        if (cvt_dims[i] == mask_set_value) mask = mask | (1 << i);
    }
    return mask;
}

void maybe_reshape_dims(dim_idx_t ndims, layout_t &layout,
        std::vector<dim_t> &dims, std::vector<dim_t> &padded_dims) {
    gpu_assert(layout.ndims() == dim_idx_t(dims.size()));
    if (layout.ndims() < ndims) {
        layout = layout_t(layout.type(), layout.blocks(), layout.offset(),
                ndims, /*do_normalize=*/false);
        dims.resize(ndims, 1);
        padded_dims.resize(ndims, 1);
    }
}

// this method only gets called when ZP precompute is in order;
// in all other cases ZPs are applied ad-hoc, without a post-op
view_t post_op_view_mapper_t::create_src_zp_view(uint32_t mask) const {
    auto map_o2k = [this](view_t &v, dim_idx_t idx, dim_t O, dim_t I, dim_t K,
                           dim_t D, dim_t P, dim_t S) {
        const auto KD = (K - 1) * (D + 1) + 1;
        const auto KDP = (KD > 1) ? KD - P : 0;
        const bool needs_right_bound = (O - 1) * S + KDP >= I;
        expr_t o = v.vvars()[idx];
        if (KD >= I) {
            o = o * S;
        } else {
            expr_t l, r;
            dim_t off = P;
            if (P > 0) l = binary_op_t::make(op_kind_t::_min, o * S - P, 0);
            if (needs_right_bound) {
                if (schedule_.var_bound(o) > O) {
                    auto q = binary_op_t::make(
                            op_kind_t::_min, o * S + KDP, (O - 1) * S + KDP);
                    r = binary_op_t::make(op_kind_t::_max, q, I);
                } else {
                    r = binary_op_t::make(op_kind_t::_max, o * S + KDP, I);
                }
                off -= I;
            }
            o = l ? l : o;
            o = r ? l ? l + r : r : o;
            o = (off != 0) ? o + off : o;
        }
        const auto &x = view_t::placeholder_var();
        auto max_pads = ir_utils::max_unique_pad_states(O, I, K, D, P, S, true);
        bool mask = max_pads
                < ir_utils::max_unique_pad_states(O, I, K, D, P, S, false);
        v.set_vdim(v.vvars()[idx], (needs_right_bound) ? O : 1);
        v.set_tdim(idx, o, (mask) ? x < max_pads : expr_t());
    };

    const auto &vars = cp_view().vvars();
    auto dst = normalize_layout(zp_dst_, /*with_groups=*/false, prb_.g,
            prb_.is_dw, prb_.dhw_map, /*add_groups=*/true, /*is_wei=*/false);
    gpu_assert((vars.size() == 6) && (dst.ndims() == 6));

    bool non_1_spatials[3] = {false, false, false};
    std::vector<layout_block_t> new_blk;
    for (auto &b : dst.blocks()) {
        if ((b.idx >= 3) && (b.idx <= 5)) {
            if (b.size > 1)
                non_1_spatials[b.idx - 3] = true;
            else if (non_1_spatials[b.idx - 3])
                continue;
        }
        new_blk.emplace_back(b);
    }
    dst = dst.with(new_blk, false);

    view_t view(vars, 6);
    view.set_vdim(vars[0], 1); // mb
    view.set_vdim(vars[1], prb_.g);
    view.set_vdim(vars[2], prb_.oc);
    view.set_tdim(0, vars[0]);
    view.set_tdim(1, vars[1]);
    view.set_tdim(2, vars[2]);
    map_o2k(view, 3, prb_.od, prb_.id, prb_.kd, prb_.dd, prb_.pd, prb_.sd);
    map_o2k(view, 4, prb_.oh, prb_.ih, prb_.kh, prb_.dh, prb_.ph, prb_.sh);
    map_o2k(view, 5, prb_.ow, prb_.iw, prb_.kw, prb_.dw, prb_.pw, prb_.sw);
    view.set_tlayout(dst);
    return view;
}

view_t post_op_view_mapper_t::create_view(const memory_desc_t &md) const {
    dim_idx_t cp_ndims = cp_view().nvdims();
    gpu_assert(cp_ndims >= 3);
    // Add groups to match ngcdhw layout.
    bool add_groups = (cp_view().vvars()[1].as<var_t>().name == "g");
    layout_t layout = make_layout(md);
    std::vector<dim_t> dims(md.dims, md.dims + md.ndims);
    std::vector<dim_t> padded_dims(md.padded_dims, md.padded_dims + md.ndims);
    maybe_reshape_dims(prb_.ndims, layout, dims, padded_dims);
    layout = normalize_layout(layout, /*with_groups=*/false, prb_.g, prb_.is_dw,
            prb_.dhw_map, add_groups,
            /*is_wei=*/false);
    auto tile = normalize_tile(dims, /*with_groups=*/false, prb_.g, prb_.is_dw,
            prb_.dhw_map, add_groups, /*is_wei=*/false);
    auto padded_tile = normalize_tile(padded_dims, /*with_groups=*/false,
            prb_.g, prb_.is_dw, prb_.dhw_map, add_groups,
            /*is_wei=*/false);
    gpu_assert(layout.ndims() == cp_ndims) << "Incompatible dimensions.";
    uint32_t bound_check_mask = 0;
    for (dim_idx_t i = 0; i < cp_ndims; i++) {
        if (tile[i] == 1) continue; // Broadcast, no bound check needed.
        if (padded_tile[i] != cp_view().tlayout().elems(i)) {
            bound_check_mask |= (1 << i);
        } else if (cp_view().has_tmask(i)) {
            bound_check_mask |= (1 << i);
        }
    }
    return view_t(layout, cp_view().vvars(), tile, bound_check_mask);
}

view_t post_op_view_mapper_t::try_create_bias_view(uint32_t mask) const {
    if ((prb_.is_fwd || prb_.is_bwd_d) && prb_.with_bias)
        return create_view(
                to_ir(prb_.conv_pd->invariant_bia_md()->data_type), mask);
    return {};
}

bool post_op_view_mapper_t::is_spurious_spatial(const pvar_t &dim) const {
    auto &var = cp_view().vvars()[dim].as<var_t>();

    int sp_idx = -1;
    if (utils::one_of(var.name, "od", "id")) {
        sp_idx = 0;
    } else if (utils::one_of(var.name, "oh", "ih")) {
        sp_idx = 1;
    } else if (utils::one_of(var.name, "ow", "iw")) {
        sp_idx = 2;
    } else {
        return false;
    }

    dim_t p = utils::pick(sp_idx, prb_.pd, prb_.ph, prb_.pw);
    dim_t s = utils::pick(sp_idx, prb_.sd, prb_.sh, prb_.sw);
    dim_t k = utils::pick(sp_idx, prb_.kd, prb_.kh, prb_.kw);
    dim_t d = utils::pick(sp_idx, prb_.dd, prb_.dh, prb_.dw);

    if (prb_.is_fwd) {
        dim_t o_value = utils::pick(sp_idx, prb_.od, prb_.oh, prb_.ow);
        dim_t o_bound = schedule_.var_bound(var);
        dim_t i = utils::pick(sp_idx, prb_.id, prb_.ih, prb_.iw);

        for (dim_t o = o_value; o < o_bound; o++) {
            dim_t i_min = o * s - p;
            if (i_min < i) return true;
        }
        return false;
    }

    if (prb_.is_bwd_d) {
        dim_t i_value = utils::pick(sp_idx, prb_.id, prb_.ih, prb_.iw);
        dim_t i_bound = schedule_.var_bound(var);
        dim_t o = utils::pick(sp_idx, prb_.od, prb_.oh, prb_.ow);

        for (dim_t i = i_value; i < i_bound; i++) {
            dim_t os_min = i - (k - 1) * (d + 1) + p;
            if (os_min < o * s) return true;
        }
        return false;
    }

    return false;
}

bool post_op_view_mapper_t::need_to_restore_zero_padding() const {
    auto &zp = prb_.conv_pd->attr()->zero_points_;
    return prb_.with_bias || !zp.has_default_values(DNNL_ARG_WEIGHTS);
}

bool post_op_view_mapper_t::use_dst_in_sum_post_op() const {
    return prb_.is_fwd;
}

bool post_op_view_mapper_t::can_use_scales() const {
    return prb_.is_fwd || prb_.is_bwd_d;
}

} // namespace jit
} // namespace conv
} // namespace intel
} // namespace gpu
} // namespace impl
} // namespace dnnl