1/*******************************************************************************
2* Copyright 2021-2022 Intel Corporation
3*
4* Licensed under the Apache License, Version 2.0 (the "License");
5* you may not use this file except in compliance with the License.
6* You may obtain a copy of the License at
7*
8* http://www.apache.org/licenses/LICENSE-2.0
9*
10* Unless required by applicable law or agreed to in writing, software
11* distributed under the License is distributed on an "AS IS" BASIS,
12* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13* See the License for the specific language governing permissions and
14* limitations under the License.
15*******************************************************************************/
16
17#ifndef GPU_JIT_CODEGEN_BANK_CONFLICT_ALLOCATION_HPP
18#define GPU_JIT_CODEGEN_BANK_CONFLICT_ALLOCATION_HPP
19
20#include "gpu/jit/codegen/reg_buf.hpp"
21#include "gpu/jit/codegen/register_allocator.hpp"
22#include "gpu/jit/ir/ir.hpp"
23#include "gpu/jit/ngen/ngen.hpp"
24
25namespace dnnl {
26namespace impl {
27namespace gpu {
28namespace jit {
29
30// Stores a set of register buffers that are allocated based on their usages
31// to avoid bank and bundle conflicts.
32class bank_conflict_allocation_t {
33public:
34 bank_conflict_allocation_t() = default;
35 bank_conflict_allocation_t(reg_allocator_t &ra) : refs_(1), ra_(&ra) {}
36
37 bool is_empty() const { return !ra_; }
38
39 int refs() const { return refs_; }
40
41 void retain() {
42 ir_assert(refs_ > 0);
43 refs_++;
44 }
45
46 void release(const expr_t &buf) {
47 ir_assert(refs_ > 0);
48 refs_--;
49 auto it = buf_map_.find(buf);
50 ir_assert(it != buf_map_.end()) << "Buffer not found: " << buf;
51 it->second.release(*ra_);
52 buf_map_.erase(it);
53 }
54
55 const reg_buf_t &get_reg_buf(const expr_t &buf) const {
56 return buf_map_.at(buf);
57 }
58
59 void set_reg_buf(const expr_t &buf, const reg_buf_t &reg_buf) {
60 auto ret = buf_map_.emplace(buf, reg_buf);
61 reg_buf.claim(*ra_);
62 ir_assert(ret.second) << "Buffer already exists: " << buf;
63 }
64
65 static bank_conflict_allocation_t create(
66 reg_allocator_t &ra, int regs, const bank_conflict_attr_t &_attr);
67
68private:
69 int refs_ = 0;
70 reg_allocator_t *ra_ = nullptr;
71 object_map_t<expr_t, reg_buf_t> buf_map_;
72};
73
74} // namespace jit
75} // namespace gpu
76} // namespace impl
77} // namespace dnnl
78
79#endif
80