blob: 0b94b19712919b2fb58506cc7315fe51bd6c0b6e [file] [log] [blame]
Tim Hall79d07d22020-04-27 18:20:16 +01001# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
Tim Hall79d07d22020-04-27 18:20:16 +010016# Description:
17# Build a live range graph for tensors in one or more subgraphs. Used for tensor allocation as well as in the scheduler.
18# Can work with either a pass packed subgraph or a scheduled subgraph.
Tim Hallffe8e282021-06-24 18:29:53 +010019from collections import namedtuple
Louis Verhaard226ecaf2021-03-30 10:18:28 +020020from typing import List
21
Tim Halld8339a72021-05-27 18:49:40 +010022import numpy as np
23
Louis Verhaardaee5d752020-09-30 09:01:52 +020024from .operation import Op
Tim Halld8339a72021-05-27 18:49:40 +010025from .tensor import MemArea
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020026from .tensor import MemType
Diego Russoe8a10452020-04-21 17:39:10 +010027from .tensor import Tensor
Tim Halld8339a72021-05-27 18:49:40 +010028from .tensor import TensorPurpose
Tim Hall79d07d22020-04-27 18:20:16 +010029
30
31class LiveRange:
Jacob Bohlin0628a8c2020-08-28 13:25:14 +020032 def __init__(self, tens, alignment):
Tim Hall79d07d22020-04-27 18:20:16 +010033 self.tensors = [] # Tensors that are assigned to the same LiveRange will be allocated to the same address
34 self.start_time = 99999999999
35 self.end_time = -1
36 self.size = 0
37 self.name = ""
Jacob Bohlin0628a8c2020-08-28 13:25:14 +020038 self.alignment = alignment
Tim Halld8339a72021-05-27 18:49:40 +010039 self.mem_area = tens.mem_area if tens else MemArea.Unknown
Tim Hall79d07d22020-04-27 18:20:16 +010040
41 if tens:
42 self.add_tensor(tens)
43
44 def __str__(self):
45 return "<live_range.LiveRange: '%s' start_time=%s, end_time=%s>" % (self.name, self.start_time, self.end_time)
46
47 __repr__ = __str__
48
49 def add_tensor(self, tens):
50 if self.size == 0:
51 self.size = tens.storage_size()
52 self.name = tens.name # LiveRange will be named after the first tensor added
53 else:
54 assert (
55 self.size >= tens.storage_size()
56 ), "Tensors assigned to the same LiveRange need to fit the size of the LiveRange."
57
58 self.tensors.append(tens)
59
Tim Halld8339a72021-05-27 18:49:40 +010060 def mark_usage(self, op_time, op_length=1):
61 op_time_start = max(op_time, 0)
62 op_time_end = op_time + op_length
63 if op_time_end <= op_time_start:
Tim Hall79d07d22020-04-27 18:20:16 +010064 return
Tim Hall79d07d22020-04-27 18:20:16 +010065
66 self.start_time = min(self.start_time, op_time_start)
67 self.end_time = max(self.end_time, op_time_end)
68
Tim Halld8339a72021-05-27 18:49:40 +010069 def set_buffer_size(self, buffer_size):
70 self.size = buffer_size
71 self.mem_area = MemArea.Sram
72
Tim Hall79d07d22020-04-27 18:20:16 +010073 def overlaps_ranges(self, other):
74 return max(self.start_time, other.start_time) < min(self.end_time, other.end_time)
75
76 def overlaps_address(self, other):
77 # Returns the first pair of tensors in this LiveRange and 'other' which have
78 # overlapping addresses
79 for tens in self.tensors:
80 for other_tens in other.tensors:
81 if max(tens.address, other_tens.address) < min(
82 tens.address + self.size, other_tens.address + other.size
83 ):
84 return True, tens, other_tens
85
86 return False, None, None
87
88 def __lt__(self, other):
89 if self.start_time != other.start_time:
90 return self.start_time < other.start_time
91 if self.end_time != other.end_time:
92 return self.end_time < other.end_time
93 if self.size != other.size:
94 return self.size < other.size
95 return self.name < other.name
96
97 def set_address(self, address):
Jacob Bohlin1a666972020-09-11 10:04:15 +020098 # Set address of all tensors in LiveRange
Tim Hall79d07d22020-04-27 18:20:16 +010099 for tens in self.tensors:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200100 tens.address = address
101
102 return address
Tim Hall79d07d22020-04-27 18:20:16 +0100103
104 def get_alignment(self):
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200105 return self.alignment
Tim Hall79d07d22020-04-27 18:20:16 +0100106
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200107 def set_alignment(self, alignment):
108 self.alignment = max(self.alignment, alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100109
110
Tim Hall79d07d22020-04-27 18:20:16 +0100111class LiveRangeGraph:
112 def __init__(self):
Louis Verhaard226ecaf2021-03-30 10:18:28 +0200113 self.lrs: List[LiveRange] = [] # List of all created ranges
Tim Hall79d07d22020-04-27 18:20:16 +0100114 self.ranges = {} # tens -> range
Tim Hall79d07d22020-04-27 18:20:16 +0100115 self.ignore_tensors = set()
116 self.processed_subgraphs = set()
117 self.current_time = 0
Tim Halld8339a72021-05-27 18:49:40 +0100118 self.end_time = None
Tim Hall79d07d22020-04-27 18:20:16 +0100119
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200120 def get_or_create_range(self, tens, alignment=Tensor.AllocationQuantum):
Jacob Bohlin1a666972020-09-11 10:04:15 +0200121 # Return the live range of the tensor (or any of its clones)
122 for existing_tensor, rng in self.ranges.items():
123 if tens.equivalent(existing_tensor):
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200124 rng.set_alignment(alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100125 return rng
126
127 # No live range found for the tensor, create a new one
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200128 rng = LiveRange(tens, alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100129 self.ranges[tens] = rng
Louis Verhaard226ecaf2021-03-30 10:18:28 +0200130 self.lrs.append(rng)
Tim Hall79d07d22020-04-27 18:20:16 +0100131 return rng
132
133 def fuse_ranges(self, in_tens, out_tens):
134 live_range = self.get_or_create_range(in_tens)
135 assert out_tens not in self.ranges, out_tens
136 live_range.add_tensor(out_tens)
137 self.ranges[out_tens] = live_range
138 return live_range
139
Tim Halld8339a72021-05-27 18:49:40 +0100140 def update_endtime(self):
Louis Verhaardcc34d5d2021-08-19 15:15:36 +0200141 self.end_time = self.current_time
Tim Halld8339a72021-05-27 18:49:40 +0100142 return self.end_time + 1
143
144 def get_temporal_memory_usage(self, target_mem_area):
Louis Verhaardcc34d5d2021-08-19 15:15:36 +0200145 usage = np.zeros(self.update_endtime(), dtype=np.int32)
Tim Halld8339a72021-05-27 18:49:40 +0100146 for rng in self.ranges.values():
147 if rng.mem_area == target_mem_area:
148 # End time is inclusive
149 usage[rng.start_time : rng.end_time + 1] += rng.size
150
151 return usage
152
Tim Hall79d07d22020-04-27 18:20:16 +0100153
Patrik Gustavssona151f592020-10-16 13:59:52 +0200154def tensor_should_be_ignored(lr_graph, tens, target_mem_area, target_mem_type_set):
155 if tens.mem_area != target_mem_area or tens.mem_type not in target_mem_type_set:
156 return True
157 if tens in lr_graph.ignore_tensors:
158 return True
Patrik Gustavssona151f592020-10-16 13:59:52 +0200159 return False
160
161
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200162def merge_elementwise_op_ranges(sched_op, lr_graph, target_mem_area, target_mem_type_set):
Tim Hallffe8e282021-06-24 18:29:53 +0100163 def _tensor_should_be_ignored(tens):
164 return tensor_should_be_ignored(lr_graph, tens, target_mem_area, target_mem_type_set)
165
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200166 # Tries to merge ifm/ofm live ranges of elementwise op
167 if sched_op.op_type.is_elementwise_op():
168 elem_op = sched_op.parent_op
Tim Hallffe8e282021-06-24 18:29:53 +0100169 if not _tensor_should_be_ignored(elem_op.ofm):
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200170 # Check if overwriting the inputs can be allowed
Tim Hallffe8e282021-06-24 18:29:53 +0100171 OpShapeTens = namedtuple("OpShapeTens", ["op_shape", "tens"])
172 outp = OpShapeTens(elem_op.ofm_shapes[0], elem_op.ofm)
173 inps = []
174 if elem_op.ifm is not None:
175 inps.append(OpShapeTens(elem_op.ifm_shapes[0], elem_op.ifm))
176 if elem_op.ifm2 is not None:
177 inps.append(OpShapeTens(elem_op.ifm_shapes[1], elem_op.ifm2))
Patrik Gustavssona151f592020-10-16 13:59:52 +0200178
Tim Hallffe8e282021-06-24 18:29:53 +0100179 # find an input tensor that can be overwritten by the output
180 for inp in inps:
181 if (
182 # check op input and output shapes allow overlapping
183 inp.op_shape == outp.op_shape
184 # check input tensor is valid
185 and inp.tens is not None
186 and inp.tens.shape != []
187 and not _tensor_should_be_ignored(inp.tens)
188 # check input and output tensors are compatible
189 and inp.tens.format == outp.tens.format
190 and inp.tens.dtype == outp.tens.dtype
191 # check input tensor only has one consumer
192 and len(inp.tens.consumer_list) == 1
193 # check output tensor only has one producer
194 and len(outp.tens.ops) == 1
195 ):
196 lr_graph.fuse_ranges(inp.tens, outp.tens)
197 break
Tim Hall79d07d22020-04-27 18:20:16 +0100198
199
200def extract_live_ranges_from_cascaded_passes(
201 sg,
202 target_mem_area,
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200203 target_mem_type_set,
Tim Hall79d07d22020-04-27 18:20:16 +0100204 ignore_subgraph_input_output_tensors=False,
205 lr_graph=None,
Tim Hallb9b515c2020-11-01 21:27:19 +0000206 cpu_tensor_alignment=Tensor.AllocationQuantum,
Tim Hall79d07d22020-04-27 18:20:16 +0100207):
Diego Russoea6111a2020-04-14 18:41:58 +0100208 if lr_graph is None:
Tim Hall79d07d22020-04-27 18:20:16 +0100209 lr_graph = LiveRangeGraph()
210
211 if sg in lr_graph.processed_subgraphs:
212 # if subgraph has been processed already, return the lr_graph as is
213 return lr_graph
214
215 if ignore_subgraph_input_output_tensors:
216 lr_graph.ignore_tensors.update(sg.input_tensors)
217 lr_graph.ignore_tensors.update(sg.output_tensors)
218
Tim Hall79d07d22020-04-27 18:20:16 +0100219 for cps in sg.cascaded_passes:
220 cps.time = lr_graph.current_time
221
222 time_for_pass = cps.time
223
Tim Hall79d07d22020-04-27 18:20:16 +0100224 for tens in cps.inputs:
Patrik Gustavssona151f592020-10-16 13:59:52 +0200225 if tensor_should_be_ignored(lr_graph, tens, target_mem_area, target_mem_type_set):
Tim Hall79d07d22020-04-27 18:20:16 +0100226 continue
Tim Hallb9b515c2020-11-01 21:27:19 +0000227 rng = lr_graph.get_or_create_range(tens, cpu_tensor_alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100228 rng.mark_usage(time_for_pass)
229
230 cps_primary_op = cps.passes[0].primary_op
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200231
Louis Verhaardaee5d752020-09-30 09:01:52 +0200232 if (
233 cps_primary_op
234 and cps_primary_op.type == Op.CustomNpuOp
235 and MemType.Permanent_CPU not in target_mem_type_set
236 ):
Tim Hall79d07d22020-04-27 18:20:16 +0100237 # If the primary-op is an NpuOp that means this is where an Npu subgraph
238 # is called. Go into said subgraph and extract live ranges before continuing.
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200239 # Use default allocation alignment of 16 for Npu tensors
Tim Hall79d07d22020-04-27 18:20:16 +0100240 npu_sg = cps_primary_op.attrs["subgraph"]
Tim Halld8339a72021-05-27 18:49:40 +0100241 lr_graph = _extract_live_ranges_from_schedule(npu_sg, target_mem_area, target_mem_type_set, lr_graph)
Tim Hall79d07d22020-04-27 18:20:16 +0100242 # Set the new time after handling the Npu subgraph
243 time_for_pass = lr_graph.current_time
244 cps.time = time_for_pass
245
Patrik Gustavssona151f592020-10-16 13:59:52 +0200246 for tens in cps.intermediates + cps.outputs:
247 if tensor_should_be_ignored(lr_graph, tens, target_mem_area, target_mem_type_set):
Tim Hall79d07d22020-04-27 18:20:16 +0100248 continue
Tim Hallb9b515c2020-11-01 21:27:19 +0000249 rng = lr_graph.get_or_create_range(tens, cpu_tensor_alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100250 rng.mark_usage(time_for_pass)
251
Tim Hall79d07d22020-04-27 18:20:16 +0100252 lr_graph.current_time += 2
253
254 end_time = 0
255 for rng in lr_graph.ranges.values():
256 # Find the maximum end time of all live-ranges in the graph
257 end_time = max(end_time, rng.end_time)
258
259 for tens in sg.output_tensors:
Patrik Gustavssona151f592020-10-16 13:59:52 +0200260 if tensor_should_be_ignored(lr_graph, tens, target_mem_area, target_mem_type_set):
Tim Hall79d07d22020-04-27 18:20:16 +0100261 continue
Tim Hallb9b515c2020-11-01 21:27:19 +0000262 rng = lr_graph.get_or_create_range(tens, cpu_tensor_alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100263 rng.mark_usage(end_time)
264
265 # Add subgraph to set of processed subgraphs
266 lr_graph.processed_subgraphs.add(sg)
267 return lr_graph
Tim Halld8339a72021-05-27 18:49:40 +0100268
269
270def create_linear_live_range_graph(sg, target_mem_area, target_mem_type_set, lr_graph):
271 assert lr_graph is not None
272 sg_time = lr_graph.current_time
273 for ps in sg.passes:
274 for tens in ps.inputs + ps.outputs + ps.intermediates:
275 if tens.purpose == TensorPurpose.Weights or tensor_should_be_ignored(
276 lr_graph, tens, target_mem_area, target_mem_type_set
277 ):
278 continue
Tim Halld8339a72021-05-27 18:49:40 +0100279 rng = lr_graph.get_or_create_range(tens)
280 rng.mark_usage(sg_time)
281
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200282 for _, op_info in sg.schedule.cost_map.items():
Tim Halld784af72021-06-08 21:25:57 +0100283 for tensor in [op_info.npu_weights_tensor, op_info.npu_scales_tensor]:
284 if tensor and not (tensor_should_be_ignored(lr_graph, tensor, target_mem_area, target_mem_type_set)):
285 rng = lr_graph.get_or_create_range(tensor)
286 rng.mark_usage(sg_time)
Tim Halld8339a72021-05-27 18:49:40 +0100287
288 lr_graph.current_time += 1
289 return lr_graph
290
291
292def _extract_live_ranges_from_schedule(sg, target_mem_area, target_mem_type_set, lr_graph):
293 time_for_cascade = {}
294 for sched_op in sg.sched_ops:
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200295 merge_elementwise_op_ranges(sched_op, lr_graph, target_mem_area, target_mem_type_set)
296
Tim Halld8339a72021-05-27 18:49:40 +0100297 op_info = sg.schedule.cost_map[sched_op]
298 cascade = op_info.cascade
299 cascade_info = sg.schedule.cascades.get(cascade, None)
300
301 time_to_set = time_for_cascade.get(cascade, lr_graph.current_time)
302
303 op_info.time_index = time_to_set
304
305 # Mark usage for all tensors related to this Pass
306 ps = sched_op.parent_ps
307 for tens in ps.inputs + ps.outputs + ps.intermediates:
308 if (
309 target_mem_area == MemArea.Sram
310 and cascade_info
311 and tens == ps.ifm_tensor
312 and sched_op in cascade_info.buffers
313 ):
314 # This tensor is a rolling buffer in a cascade and the size of the LiveRange needs to be modified
315 # for enabling temporal memory snapshots without modifying the original Tensor
316 rng = lr_graph.get_or_create_range(tens)
317 rng.set_buffer_size(cascade_info.buffers[sched_op].elements() * sched_op.ifm.dtype.size_in_bytes())
318 elif (
319 tens.purpose == TensorPurpose.Weights
320 or tens.purpose == TensorPurpose.FSBias
321 or tens.mem_type not in target_mem_type_set
322 or tens.mem_area != target_mem_area
323 ):
324 continue
325
326 else:
327 rng = lr_graph.get_or_create_range(tens)
328
329 rng.mark_usage(time_to_set)
330
331 weight_tens = op_info.buffered_weight_tensor
332 if weight_tens and weight_tens.mem_type in target_mem_type_set and weight_tens.mem_area == target_mem_area:
333 rng = lr_graph.get_or_create_range(weight_tens)
334 if weight_tens.pre_buffer:
335 rng.mark_usage(time_to_set - 1, 2)
336 else:
337 rng.mark_usage(time_to_set)
338
339 if time_to_set == lr_graph.current_time:
340 lr_graph.current_time += 2
341
342 if cascade != 0:
343 time_for_cascade[cascade] = time_to_set
344
345 end_time = lr_graph.update_endtime()
346
347 for tens in sg.output_tensors:
348 if tens.mem_type not in target_mem_type_set or tens.mem_area != target_mem_area:
349 continue
350 rng = lr_graph.get_or_create_range(tens)
351 rng.mark_usage(end_time)
352
353 return lr_graph