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_REGISTER_ALLOCATOR_HPP |
18 | #define GPU_JIT_CODEGEN_REGISTER_ALLOCATOR_HPP |
19 | |
20 | #include "common/z_magic.hpp" |
21 | #include "gpu/jit/ngen/ngen.hpp" |
22 | #include "gpu/jit/ngen/ngen_register_allocator.hpp" |
23 | #include "gpu/jit/utils/utils.hpp" |
24 | |
25 | namespace dnnl { |
26 | namespace impl { |
27 | namespace gpu { |
28 | namespace jit { |
29 | |
30 | // Register Allocator Wrapper to allow for custom checks. |
31 | class reg_allocator_t { |
32 | public: |
33 | static const uint64_t warn_none = 0; |
34 | static const uint64_t warn_large_grf = 1; |
35 | static const uint64_t warn_all = warn_large_grf; |
36 | static const uint64_t warn_default = warn_none; |
37 | |
38 | // Based on nGEN limitation where subregisters are allocated in dword chunks. |
39 | static const int granularity = 4; |
40 | |
41 | reg_allocator_t(ngen::HW hw, const std::string &kernel_name_, |
42 | int warn_flags_ = warn_default) |
43 | : ra(hw) { |
44 | #ifdef GEN_CONV_DEBUG |
45 | kernel_name = kernel_name_; |
46 | warn_flags = warn_flags_; |
47 | #endif |
48 | MAYBE_UNUSED(kernel_name_); |
49 | MAYBE_UNUSED(warn_flags_); |
50 | } |
51 | ~reg_allocator_t() { |
52 | #ifdef GEN_CONV_DEBUG |
53 | if ((warn_flags & warn_large_grf) && (peak_grf_usage <= 128) |
54 | && (ra.getRegisterCount() > 128)) |
55 | ir_warning() << kernel_name |
56 | << " uselessly enables large grf mode as " |
57 | << peak_grf_usage << " registers were used\n" ; |
58 | ir_assert(!is_speculate) << "Speculative allocation never finished\n" ; |
59 | #endif |
60 | } |
61 | |
62 | ngen::HW hardware() const { return ra.hardware(); } |
63 | |
64 | ngen::GRFRange alloc_range(int nregs, |
65 | ngen::Bundle base_bundle = ngen::Bundle(), |
66 | ngen::BundleGroup bundle_mask = ngen::BundleGroup::AllBundles()) { |
67 | auto ret = ra.alloc_range(nregs, base_bundle, bundle_mask); |
68 | update_peak_grf_usage(); |
69 | return ret; |
70 | } |
71 | ngen::GRF alloc(ngen::Bundle bundle = ngen::Bundle()) { |
72 | auto ret = ra.alloc(bundle); |
73 | update_peak_grf_usage(); |
74 | return ret; |
75 | } |
76 | |
77 | ngen::FlagRegister alloc_flag() { return ra.alloc_flag(); } |
78 | |
79 | ngen::GRFRange try_alloc_range(int nregs, |
80 | ngen::Bundle base_bundle = ngen::Bundle(), |
81 | ngen::BundleGroup bundle_mask = ngen::BundleGroup::AllBundles()) { |
82 | auto ret = ra.try_alloc_range(nregs, base_bundle, bundle_mask); |
83 | update_peak_grf_usage(); |
84 | return ret; |
85 | } |
86 | ngen::GRF try_alloc(ngen::Bundle bundle = ngen::Bundle()) { |
87 | auto ret = ra.try_alloc(bundle); |
88 | update_peak_grf_usage(); |
89 | return ret; |
90 | } |
91 | |
92 | ngen::Subregister alloc_sub( |
93 | ngen::DataType type, ngen::Bundle bundle = ngen::Bundle()) { |
94 | auto ret = ra.alloc_sub(type, bundle); |
95 | update_peak_grf_usage(); |
96 | return ret; |
97 | } |
98 | |
99 | template <typename T> |
100 | ngen::Subregister alloc_sub(ngen::Bundle bundle = ngen::Bundle()) { |
101 | auto ret = ra.alloc_sub<T>(bundle); |
102 | update_peak_grf_usage(); |
103 | return ret; |
104 | } |
105 | |
106 | ngen::Subregister try_alloc_sub( |
107 | ngen::DataType type, ngen::Bundle bundle = ngen::Bundle()) { |
108 | auto ret = ra.try_alloc_sub(type, bundle); |
109 | update_peak_grf_usage(); |
110 | return ret; |
111 | } |
112 | template <typename T> |
113 | ngen::Subregister try_alloc_sub(ngen::Bundle bundle = ngen::Bundle()) { |
114 | auto ret = ra.try_alloc_sub<T>(bundle); |
115 | update_peak_grf_usage(); |
116 | return ret; |
117 | } |
118 | template <typename RD> |
119 | void safeRelease(RD ®) { |
120 | ra.safeRelease(reg); |
121 | } |
122 | template <typename RD> |
123 | void release(RD reg) { |
124 | ra.release(reg); |
125 | } |
126 | template <typename RD> |
127 | void claim(RD reg) { |
128 | ra.claim(reg); |
129 | update_peak_grf_usage(); |
130 | } |
131 | |
132 | void setRegisterCount(int rcount) { ra.setRegisterCount(rcount); } |
133 | |
134 | #ifdef GEN_CONV_DEBUG |
135 | int get_peak_grf_usage() const { return peak_grf_usage; } |
136 | int get_grf_usage() const { return ra.countAllocedRegisters(); } |
137 | |
138 | // For performing speculative allocations that may not be used in the final |
139 | // register allocation |
140 | void start_speculate() { is_speculate = true; } |
141 | void finish_speculate() { |
142 | is_speculate = false; |
143 | update_peak_grf_usage(); |
144 | } |
145 | #else |
146 | void start_speculate() {} |
147 | void finish_speculate() {} |
148 | #endif |
149 | |
150 | protected: |
151 | #ifdef GEN_CONV_DEBUG |
152 | void update_peak_grf_usage() { |
153 | if (is_speculate) return; |
154 | int register_count = get_grf_usage(); |
155 | if (peak_grf_usage < register_count) peak_grf_usage = register_count; |
156 | } |
157 | #else |
158 | void update_peak_grf_usage() {} |
159 | #endif |
160 | |
161 | int peak_grf_usage = 0; |
162 | int warn_flags; |
163 | bool is_speculate = false; |
164 | std::string kernel_name; |
165 | |
166 | private: |
167 | ngen::RegisterAllocator ra; |
168 | }; |
169 | |
170 | } // namespace jit |
171 | } // namespace gpu |
172 | } // namespace impl |
173 | } // namespace dnnl |
174 | #endif |
175 | |