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
25namespace dnnl {
26namespace impl {
27namespace gpu {
28namespace jit {
29
30// Register Allocator Wrapper to allow for custom checks.
31class reg_allocator_t {
32public:
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 &reg) {
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
150protected:
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
166private:
167 ngen::RegisterAllocator ra;
168};
169
170} // namespace jit
171} // namespace gpu
172} // namespace impl
173} // namespace dnnl
174#endif
175