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/jit/pass/bank_conflict.hpp"

#include "gemmstone/../../dsl/ir/pass/trace.hpp"
#include "gpu/intel/jit/ir/legacy.hpp"

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

// FIXME: Use convolution-agnostic mechanism to skip zero-points related calls.
static bool is_zero_points_call(const stmt_t &s) {
    auto is_zp_var = [&](const expr_t &e) {
        auto &base = get_base(e);
        auto &name = base.as<var_t>().name;
        return name.find("zp_") == 0;
    };
    if (is_func_call<dpas_t>(s)) {
        auto &src1 = dpas_t::arg_src1(s);
        auto &src2 = dpas_t::arg_src2(s);
        return is_zp_var(src1) || is_zp_var(src2);
    }
    if (is_func_call<mad_t>(s)) {
        auto &src1 = mad_t::arg_src1(s);
        auto &src2 = mad_t::arg_src2(s);
        return is_zp_var(src1) || is_zp_var(src2);
    }
    return false;
}

class bank_conflict_attribute_injector_t : public ir_mutator_t {
public:
    object_t _mutate(const alloc_t &obj) override {
        all_buf_sizes_.emplace(obj.buf, obj.size);

        auto new_obj = ir_mutator_t::_mutate(obj);
        if (bufs_.count(obj.buf) == 0) return new_obj;

        init_attr();

        auto new_attrs = obj.attrs;
        new_attrs.push_back(attr_);
        auto &body = new_obj.as<alloc_t>().body;
        return alloc_t::make(obj.buf, obj.size, obj.kind, new_attrs, body);
    }

    object_t _mutate(const func_call_t &obj) override {
        if (is_frozen) return ir_mutator_t::_mutate(obj);
        if (is_zero_points_call(obj)) return ir_mutator_t::_mutate(obj);

        bool is_mad = obj.func.is<mad_t>();
        bool is_dpas = obj.func.is<dpas_t>();
        auto *send = obj.func.as_ptr<send_t>();
        bool is_load = send && (send->is_load() || send->is_load_2d());

        if (is_mad || is_dpas) {
            auto dst_buf = ptr_base(obj.args[0]);
            auto src0_buf = ptr_base(obj.args[1]);
            auto src1_buf = ptr_base(obj.args[2]);
            auto src2_buf = ptr_base(obj.args[3]);

            // src0 may be null in some cases, skip it.
            if (src0_buf) bufs_.insert(std::move(src0_buf));
            bufs_.insert(std::move(src1_buf));
            bufs_.insert(std::move(src2_buf));

            instructions_.emplace_back(obj);
        } else if (is_load) {
            auto &buf = send_t::arg_reg_buf(obj);
            auto &base = (is_var(buf) ? buf : buf.as<ptr_t>().base);
            int off = (is_var(buf) ? 0 : to_cpp<int>(buf.as<ptr_t>().off));
            int size = send->payload_size();
            auto span = [](int a, int b) {
                // Returns minimal 2^B so that there is x such that:
                //   x * 2^B <= a <= b < (x + 1) * 2^B
                int mask = a ^ b;
                // Neither of the highest two bits may be set
                //  - 32nd bit set: a is negative, b is nonnegative, 2^B does
                //                  not exist for integer x
                //  - 31st bit set: requested span is 0x80000000 < 0 in
                //                  integer precision
                gpu_assert(!(mask & 0xc0000000))
                        << "span cannot be represented in integer precision";
                return utils::rnd_up_pow2(mask + 1);
            }(off, off + size - 1);
            int &min_block_size = all_buf_min_block_sizes[base];
            min_block_size = std::max(min_block_size, span);
        }
        return ir_mutator_t::_mutate(obj);
    }

private:
    void init_attr() {
        if (attr_) return;

        is_frozen = true;
        std::vector<expr_t> buf_vec;
        std::vector<int> buf_sizes;
        std::vector<int> buf_min_block_sizes;
        for (auto &buf : bufs_) {
            buf_vec.push_back(buf);
            buf_sizes.push_back(all_buf_sizes_.at(buf));
            auto it = all_buf_min_block_sizes.find(buf);
            int min_block_size
                    = (it == all_buf_min_block_sizes.end() ? 0 : it->second);
            buf_min_block_sizes.push_back(min_block_size);
        }
        attr_ = bank_conflict_attr_t::make(
                buf_vec, buf_sizes, buf_min_block_sizes, instructions_);
    }

    static expr_t ptr_base(const expr_t &e) {
        if (e.is<var_t>()) return e;
        auto *ptr = e.as_ptr<ptr_t>();
        if (ptr) return e.as<ptr_t>().base;
        return expr_t();
    }

    object_map_t<expr_t, int> all_buf_sizes_;
    object_map_t<expr_t, int> all_buf_min_block_sizes;
    object_eq_set_t<expr_t> bufs_;
    std::vector<stmt_t> instructions_;
    bool is_frozen = false;

    alloc_attr_t attr_;
};

stmt_t inject_bank_conflict_attribute(const stmt_t &s, ir_context_t &ir_ctx) {
    ir::trace_start();
    auto ret = bank_conflict_attribute_injector_t().mutate(s);
    ir::trace_pass("inject_bank_conflict_attribute", ret, ir_ctx);
    return ret;
}

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