onednn-src 0.1.13

/*******************************************************************************
* Copyright 2021 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 <math.h>

#include "common/primitive_exec_types.hpp"

#include "gpu/intel/concat/xe.hpp"

namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace concat {

std::pair<int, int> xe_t::pd_t::calculate_iter_dim_idx_chunk(
        int num_threads) const {
    if (conf.ndims <= 1) return std::make_pair(0, 1);
    const auto &dst_dims = conf.dst_md_info.padded_dims;
    int max_dim_idx = (conf.concat_axis == conf.ndims - 1) ? conf.ndims - 2
                                                           : conf.ndims - 1;
    dim_t max_dim = dst_dims[max_dim_idx];
    for (int dim_idx = max_dim_idx - 1; dim_idx >= 0; dim_idx--) {
        if (dim_idx == conf.concat_axis) continue;
        const auto dim = dst_dims[dim_idx];
        if (dim > max_dim) {
            max_dim = dim;
            max_dim_idx = dim_idx;
        }
    }
    const int iter_dim_idx = max_dim_idx;
    const dim_t all_elems = utils::array_product(dst_dims, conf.ndims);
    const int max_iter_dim_chunk = 1024;
    const int min_threads = num_threads * 4;
    int iter_dim_chunk = static_cast<int>(
            std::min<dim_t>(dst_dims[iter_dim_idx], max_iter_dim_chunk));
    const auto get_num_threads = [&]() {
        return utils::div_up(all_elems, iter_dim_chunk * conf.sub_group_size);
    };
    while (get_num_threads() < min_threads && iter_dim_chunk > 1) {
        iter_dim_chunk = utils::div_up(iter_dim_chunk, 2);
    }
    return std::make_pair(iter_dim_idx, iter_dim_chunk);
}

bool xe_t::pd_t::can_use_sub_group_size(
        const intel::engine_t *intel_engine, int sub_group_size) const {
    auto is_dim_dense = [](const memory_desc_wrapper &mdw, int dim_idx) {
        return mdw.blocking_desc().strides[dim_idx] == 1;
    };
    auto get_dim_block = [](const memory_desc_wrapper &mdw, int dim_idx) {
        const auto &blk = mdw.blocking_desc();
        if (blk.inner_nblks == 0
                || blk.inner_idxs[blk.inner_nblks - 1] != dim_idx)
            return static_cast<dnnl_dim_t>(1);
        return blk.inner_blks[blk.inner_nblks - 1];
    };

    const int c_idx = 1;
    const memory_desc_wrapper dst_mdw(dst_md());
    const bool is_dst_blocked = dst_mdw.blocking_desc().inner_nblks > 0;
    const int max_sub_group_size
            = intel_engine->device_info()->max_subgroup_size();
    bool is_concat_axis_aligned = true;
    bool layouts_compatible = is_dst_blocked
            ? get_dim_block(dst_mdw, c_idx) % sub_group_size == 0
            : is_dim_dense(dst_mdw, c_idx);
    for (int i = 0; i < conf.n; ++i) {
        const memory_desc_wrapper src_mdw(src_md(i));
        is_concat_axis_aligned = is_concat_axis_aligned
                && src_mdw.dims()[conf.concat_axis] % sub_group_size == 0;
        if (is_dst_blocked) {
            layouts_compatible = layouts_compatible
                    && get_dim_block(src_mdw, c_idx) % sub_group_size == 0;
        } else {
            layouts_compatible
                    = layouts_compatible && is_dim_dense(src_mdw, c_idx);
        }
    }
    return is_concat_axis_aligned && layouts_compatible
            && sub_group_size <= max_sub_group_size
            && intel_engine->mayiuse_sub_group(sub_group_size);
}

int xe_t::pd_t::calculate_sub_group_size(
        const intel::engine_t *intel_engine) const {
    // Subgroups are used only for concatenation over C dimension
    if (conf.concat_axis != 1) return 1;
    for (int sub_group_size : {32, 16, 8}) {
        if (can_use_sub_group_size(intel_engine, sub_group_size)) {
            return sub_group_size;
        }
    }
    return 1;
}

status_t xe_t::pd_t::init_conf(impl::engine_t *engine) {
    const concat_pd_t *pd = this;

    const memory_desc_wrapper dst_mdw(pd->dst_md());
    const auto dst_dims = dst_mdw.padded_dims();

    conf.dst_md_info = memory_desc_info_t::create(dst_mdw);
    conf.dst_type = dst_mdw.data_type();
    conf.dst_offset0 = dst_mdw.offset0();
    conf.src_type = memory_desc_wrapper(pd->src_md(0)).data_type();
    conf.ndims = dst_mdw.ndims();
    const auto *intel_engine = utils::downcast<intel::engine_t *>(engine);
    conf.dispatch = intel_engine->create_dispatch(dst_mdw.md_);
    conf.n = pd->n_inputs();
    conf.concat_axis = pd->concat_dim();

    dim_t max_elems = dst_mdw.nelems();
    dim_t concat_axis_end = 0;
    conf.scales_mask = 0;
    for (int i = 0; i < conf.n; ++i) {
        const memory_desc_wrapper src_mdw(pd->src_md(i));
        max_elems = std::max(max_elems, into<dim_t>(src_mdw.nelems()));
        concat_axis_end += src_mdw.dims()[conf.concat_axis];
        conf.offset[i] = concat_axis_end;
        conf.src_md_infos[i] = memory_desc_info_t::create(pd->src_md(i));
        conf.scale_src[i] = scales_query_t(
                pd->attr(), src_mdw, DNNL_ARG_MULTIPLE_SRC + i);
        conf.scales_mask |= ((!conf.scale_src[i].has_default_values()) << i);
    }

    conf.use_large_index = (max_elems > std::numeric_limits<int>::max());
    conf.require_stateless_addressing = pd->has_large_buffers();

    conf.sub_group_size = calculate_sub_group_size(intel_engine);
    std::tie(conf.iter_dim_idx, conf.iter_dim_chunk)
            = calculate_iter_dim_idx_chunk(
                    intel_engine->device_info()->hw_threads());

    int sg_chunk = conf.iter_dim_chunk * conf.sub_group_size;
    dim_t chunks = dst_mdw.nelems(true) / sg_chunk;
    VDISPATCH_CONCAT_IC(chunks <= (1u << 15) || conf.sub_group_size != 1,
            VERBOSE_BAD_PARAM, "sub_group_size");
    VDISPATCH_CONCAT_IC(dst_mdw.blocking_desc().inner_nblks != 0
                    || !(conf.sub_group_size == 1
                            || (conf.ndims > 2 && conf.iter_dim_chunk == 1)),
            VERBOSE_BAD_PARAM, "ndims, dim_chunk for non-unit group size");

    for (int dim_idx = 0; dim_idx < MAX_NDIMS; dim_idx++) {
        const int dim_block
                = conf.iter_dim_idx == dim_idx ? conf.iter_dim_chunk : 1;
        const dim_t dim_size = conf.ndims > dim_idx ? dst_dims[dim_idx] : 1;
        conf.dispatch.define_dim(
                utils::format("D%d", dim_idx), 0, dim_size, dim_block);
    }
    if (conf.sub_group_size > 1) {
        CHECK(conf.dispatch.vectorize_dim("D1", conf.sub_group_size));
    }
    conf.dispatch.generate();

    return status::success;
}

static status_t init_kernel_ctx_common(
        compute::kernel_ctx_t &kernel_ctx, const conf_t &conf) {
    constexpr bool with_punning = false;

    for (int i = 0; i < conf.n; ++i) {
        kernel_ctx.define_int(utils::format("SRC%d_END", i), conf.offset[i]);
        def_memory_desc_info(kernel_ctx, conf.src_md_infos[i],
                utils::format("SRC%d", i).c_str(), with_punning);
    }
    def_memory_desc_info(kernel_ctx, conf.dst_md_info, "DST", with_punning);

    kernel_ctx.set_data_type(conf.src_type, with_punning);
    kernel_ctx.require_stateless_addressing(conf.require_stateless_addressing);

    kernel_ctx.define_int("NDIMS", conf.ndims);
    kernel_ctx.define_int("CONCAT_AXIS", conf.concat_axis);
    kernel_ctx.define_int("NUM_INPUTS", conf.n);
    kernel_ctx.define_int("SUB_GROUP_SIZE", conf.sub_group_size);
    kernel_ctx.define_int("VECT_DT_N", 1);
    kernel_ctx.define_int("ITER_DIM_PADDED_SIZE",
            conf.dst_md_info.padded_dims[conf.iter_dim_idx]);
    kernel_ctx.define_int("ITER_DIM_IDX", conf.iter_dim_idx);
    kernel_ctx.define_int("ITER_DIM_CHUNK", conf.iter_dim_chunk);
    kernel_ctx.define_int("SCALES_MASK", conf.scales_mask);

    kernel_ctx.use_int32_offset(!conf.use_large_index);

    def_dispatch(kernel_ctx, conf.dispatch);

    return status::success;
}

status_t xe_t::pd_t::init_kernel_ctx(compute::kernel_ctx_t &kernel_ctx) const {
    return init_kernel_ctx_common(kernel_ctx, conf);
}

status_t xe_t::execute(const exec_ctx_t &ctx) const {
    auto &dst = CTX_OUT_STORAGE(DNNL_ARG_DST);
    int next_arg = 0;
    const auto &conf = pd()->conf;
    compute::kernel_arg_list_t arg_list;
    arg_list.set(next_arg++, dst);
    arg_list.set(next_arg++, conf.dst_offset0);
    for (int i = 0; i < 16; ++i) {
        auto &src = CTX_IN_STORAGE(DNNL_ARG_MULTIPLE_SRC + i);
        arg_list.set(next_arg++, src);
    }

    // add scales for all inputs
    if (conf.scales_mask) {
        for (int i = 0; i < pd()->n_inputs(); ++i) {
            arg_list.set(next_arg++, conf.scale_src[i].get_scales(ctx));
        }
    }

    auto nd_range = conf.dispatch.nd_range();

    return parallel_for(ctx, nd_range, kernel, arg_list);
}

} // namespace concat
} // namespace intel
} // namespace gpu
} // namespace impl
} // namespace dnnl