blob: b18afeccfd5c0aa7da85f298a85c333d2b4531c9 [file] [log] [blame]
Fredrik Svedberg33c01e62023-02-13 11:32:12 +01001# SPDX-FileCopyrightText: Copyright 2020-2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
Tim Hall79d07d22020-04-27 18:20:16 +01002#
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.
Rickard Bolinbc6ee582022-11-04 08:24:29 +000016#
Tim Hall79d07d22020-04-27 18:20:16 +010017# Description:
18# Build a live range graph for tensors in one or more subgraphs. Used for tensor allocation as well as in the scheduler.
19# Can work with either a pass packed subgraph or a scheduled subgraph.
Tim Hallffe8e282021-06-24 18:29:53 +010020from collections import namedtuple
Louis Verhaard226ecaf2021-03-30 10:18:28 +020021from typing import List
22
Tim Halld8339a72021-05-27 18:49:40 +010023import numpy as np
24
Louis Verhaardaee5d752020-09-30 09:01:52 +020025from .operation import Op
Tim Halld8339a72021-05-27 18:49:40 +010026from .tensor import MemArea
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020027from .tensor import MemType
Diego Russoe8a10452020-04-21 17:39:10 +010028from .tensor import Tensor
Tim Halld8339a72021-05-27 18:49:40 +010029from .tensor import TensorPurpose
Tim Hall79d07d22020-04-27 18:20:16 +010030
31
32class LiveRange:
Jacob Bohlin0628a8c2020-08-28 13:25:14 +020033 def __init__(self, tens, alignment):
Tim Hall79d07d22020-04-27 18:20:16 +010034 self.tensors = [] # Tensors that are assigned to the same LiveRange will be allocated to the same address
35 self.start_time = 99999999999
36 self.end_time = -1
37 self.size = 0
38 self.name = ""
Jacob Bohlin0628a8c2020-08-28 13:25:14 +020039 self.alignment = alignment
Tim Halld8339a72021-05-27 18:49:40 +010040 self.mem_area = tens.mem_area if tens else MemArea.Unknown
Tim Hall79d07d22020-04-27 18:20:16 +010041
42 if tens:
43 self.add_tensor(tens)
44
45 def __str__(self):
erik.andersson@arm.comde6cb642022-02-02 14:03:15 +010046 return (
47 f"<live_range.LiveRange: {self.start_time}-{self.end_time}, "
48 f"size={self.size}, '{self.name}' #:{len(self.tensors)}>"
49 )
Tim Hall79d07d22020-04-27 18:20:16 +010050
51 __repr__ = __str__
52
53 def add_tensor(self, tens):
54 if self.size == 0:
55 self.size = tens.storage_size()
56 self.name = tens.name # LiveRange will be named after the first tensor added
57 else:
58 assert (
59 self.size >= tens.storage_size()
60 ), "Tensors assigned to the same LiveRange need to fit the size of the LiveRange."
61
62 self.tensors.append(tens)
63
Tim Halld8339a72021-05-27 18:49:40 +010064 def mark_usage(self, op_time, op_length=1):
65 op_time_start = max(op_time, 0)
66 op_time_end = op_time + op_length
Rickard Bolinfd8b5002022-05-16 09:11:06 +000067 if op_time_end < op_time_start:
Tim Hall79d07d22020-04-27 18:20:16 +010068 return
Tim Hall79d07d22020-04-27 18:20:16 +010069
70 self.start_time = min(self.start_time, op_time_start)
71 self.end_time = max(self.end_time, op_time_end)
72
Tim Halld8339a72021-05-27 18:49:40 +010073 def set_buffer_size(self, buffer_size):
74 self.size = buffer_size
75 self.mem_area = MemArea.Sram
76
Tim Hall79d07d22020-04-27 18:20:16 +010077 def overlaps_ranges(self, other):
78 return max(self.start_time, other.start_time) < min(self.end_time, other.end_time)
79
80 def overlaps_address(self, other):
81 # Returns the first pair of tensors in this LiveRange and 'other' which have
82 # overlapping addresses
83 for tens in self.tensors:
84 for other_tens in other.tensors:
85 if max(tens.address, other_tens.address) < min(
86 tens.address + self.size, other_tens.address + other.size
87 ):
88 return True, tens, other_tens
89
90 return False, None, None
91
92 def __lt__(self, other):
93 if self.start_time != other.start_time:
94 return self.start_time < other.start_time
95 if self.end_time != other.end_time:
96 return self.end_time < other.end_time
97 if self.size != other.size:
98 return self.size < other.size
99 return self.name < other.name
100
101 def set_address(self, address):
Jacob Bohlin1a666972020-09-11 10:04:15 +0200102 # Set address of all tensors in LiveRange
Tim Hall79d07d22020-04-27 18:20:16 +0100103 for tens in self.tensors:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200104 tens.address = address
105
106 return address
Tim Hall79d07d22020-04-27 18:20:16 +0100107
108 def get_alignment(self):
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200109 return self.alignment
Tim Hall79d07d22020-04-27 18:20:16 +0100110
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200111 def set_alignment(self, alignment):
112 self.alignment = max(self.alignment, alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100113
114
Tim Hall79d07d22020-04-27 18:20:16 +0100115class LiveRangeGraph:
116 def __init__(self):
Louis Verhaard226ecaf2021-03-30 10:18:28 +0200117 self.lrs: List[LiveRange] = [] # List of all created ranges
Tim Hall79d07d22020-04-27 18:20:16 +0100118 self.ranges = {} # tens -> range
Tim Hall79d07d22020-04-27 18:20:16 +0100119 self.processed_subgraphs = set()
120 self.current_time = 0
Tim Halld8339a72021-05-27 18:49:40 +0100121 self.end_time = None
Tim Hall79d07d22020-04-27 18:20:16 +0100122
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200123 def get_or_create_range(self, tens, alignment=Tensor.AllocationQuantum):
Jacob Bohlin1a666972020-09-11 10:04:15 +0200124 # Return the live range of the tensor (or any of its clones)
125 for existing_tensor, rng in self.ranges.items():
126 if tens.equivalent(existing_tensor):
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200127 rng.set_alignment(alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100128 return rng
129
130 # No live range found for the tensor, create a new one
Jacob Bohlin0628a8c2020-08-28 13:25:14 +0200131 rng = LiveRange(tens, alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100132 self.ranges[tens] = rng
Louis Verhaard226ecaf2021-03-30 10:18:28 +0200133 self.lrs.append(rng)
Tim Hall79d07d22020-04-27 18:20:16 +0100134 return rng
135
136 def fuse_ranges(self, in_tens, out_tens):
137 live_range = self.get_or_create_range(in_tens)
138 assert out_tens not in self.ranges, out_tens
139 live_range.add_tensor(out_tens)
140 self.ranges[out_tens] = live_range
141 return live_range
142
Tim Halld8339a72021-05-27 18:49:40 +0100143 def update_endtime(self):
Louis Verhaardcc34d5d2021-08-19 15:15:36 +0200144 self.end_time = self.current_time
Tim Halld8339a72021-05-27 18:49:40 +0100145 return self.end_time + 1
146
147 def get_temporal_memory_usage(self, target_mem_area):
Louis Verhaardcc34d5d2021-08-19 15:15:36 +0200148 usage = np.zeros(self.update_endtime(), dtype=np.int32)
erik.andersson@arm.comde6cb642022-02-02 14:03:15 +0100149 for lr in self.lrs:
150 if lr.mem_area == target_mem_area:
Tim Halld8339a72021-05-27 18:49:40 +0100151 # End time is inclusive
erik.andersson@arm.comde6cb642022-02-02 14:03:15 +0100152 usage[lr.start_time : lr.end_time + 1] += lr.size
Tim Halld8339a72021-05-27 18:49:40 +0100153
154 return usage
155
Tim Hall79d07d22020-04-27 18:20:16 +0100156
Fredrik Svedberg0ae28482021-10-27 13:58:03 +0200157def tensor_should_be_ignored(tens, target_mem_area, target_mem_type_set):
Johan Alfvénfba0a7d2022-10-11 20:41:41 +0200158 if target_mem_area is None or target_mem_type_set is None:
159 return False
Patrik Gustavssona151f592020-10-16 13:59:52 +0200160 if tens.mem_area != target_mem_area or tens.mem_type not in target_mem_type_set:
161 return True
Patrik Gustavssona151f592020-10-16 13:59:52 +0200162 return False
163
164
Johan Alfvénfba0a7d2022-10-11 20:41:41 +0200165def _get_ifm_to_fuse(sched_op, target_mem_area=None, target_mem_type_set=None):
Tim Hallffe8e282021-06-24 18:29:53 +0100166 def _tensor_should_be_ignored(tens):
Johan Alfvén8d57aaa2022-02-04 11:19:17 +0100167 if tens.ifm_write_protected:
Fredrik Svedberg0ae28482021-10-27 13:58:03 +0200168 return True
169 return tensor_should_be_ignored(tens, target_mem_area, target_mem_type_set)
Tim Hallffe8e282021-06-24 18:29:53 +0100170
Johan Alfvénfba0a7d2022-10-11 20:41:41 +0200171 # Check if possible to merge ifm/ofm live ranges of elementwise op
172 ifm_tens = None
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200173 if sched_op.op_type.is_elementwise_op():
174 elem_op = sched_op.parent_op
Tim Hallffe8e282021-06-24 18:29:53 +0100175 if not _tensor_should_be_ignored(elem_op.ofm):
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200176 # Check if overwriting the inputs can be allowed
Tim Hallffe8e282021-06-24 18:29:53 +0100177 OpShapeTens = namedtuple("OpShapeTens", ["op_shape", "tens"])
178 outp = OpShapeTens(elem_op.ofm_shapes[0], elem_op.ofm)
179 inps = []
180 if elem_op.ifm is not None:
181 inps.append(OpShapeTens(elem_op.ifm_shapes[0], elem_op.ifm))
182 if elem_op.ifm2 is not None:
183 inps.append(OpShapeTens(elem_op.ifm_shapes[1], elem_op.ifm2))
Patrik Gustavssona151f592020-10-16 13:59:52 +0200184
Tim Hallffe8e282021-06-24 18:29:53 +0100185 # find an input tensor that can be overwritten by the output
186 for inp in inps:
187 if (
188 # check op input and output shapes allow overlapping
189 inp.op_shape == outp.op_shape
190 # check input tensor is valid
191 and inp.tens is not None
192 and inp.tens.shape != []
193 and not _tensor_should_be_ignored(inp.tens)
194 # check input and output tensors are compatible
195 and inp.tens.format == outp.tens.format
196 and inp.tens.dtype == outp.tens.dtype
197 # check input tensor only has one consumer
198 and len(inp.tens.consumer_list) == 1
199 # check output tensor only has one producer
200 and len(outp.tens.ops) == 1
201 ):
Johan Alfvénfba0a7d2022-10-11 20:41:41 +0200202 ifm_tens = inp.tens
Tim Hallffe8e282021-06-24 18:29:53 +0100203 break
Tim Hall79d07d22020-04-27 18:20:16 +0100204
Johan Alfvénfba0a7d2022-10-11 20:41:41 +0200205 return ifm_tens
206
207
208def ofm_can_reuse_ifm(sched_op, target_mem_area=None, target_mem_type_set=None):
209 ifm = _get_ifm_to_fuse(sched_op, target_mem_area, target_mem_type_set)
210 return ifm is not None
211
212
213def merge_elementwise_op_ranges(sg, sched_op, lr_graph, target_mem_area, target_mem_type_set):
214 ifm = _get_ifm_to_fuse(sched_op, target_mem_area, target_mem_type_set)
215 if ifm:
216 lr_graph.fuse_ranges(ifm, sched_op.parent_op.ofm)
217
Tim Hall79d07d22020-04-27 18:20:16 +0100218
219def extract_live_ranges_from_cascaded_passes(
Jonas Ohlssond8575072022-03-30 10:30:25 +0200220 sg,
221 target_mem_area,
222 target_mem_type_set,
223 lr_graph=None,
224 cpu_tensor_alignment=Tensor.AllocationQuantum,
Tim Hall79d07d22020-04-27 18:20:16 +0100225):
Diego Russoea6111a2020-04-14 18:41:58 +0100226 if lr_graph is None:
Tim Hall79d07d22020-04-27 18:20:16 +0100227 lr_graph = LiveRangeGraph()
228
229 if sg in lr_graph.processed_subgraphs:
230 # if subgraph has been processed already, return the lr_graph as is
231 return lr_graph
232
Tim Hall79d07d22020-04-27 18:20:16 +0100233 for cps in sg.cascaded_passes:
234 cps.time = lr_graph.current_time
235
236 time_for_pass = cps.time
237
Tim Hall79d07d22020-04-27 18:20:16 +0100238 for tens in cps.inputs:
Fredrik Svedberg0ae28482021-10-27 13:58:03 +0200239 if tensor_should_be_ignored(tens, target_mem_area, target_mem_type_set):
Tim Hall79d07d22020-04-27 18:20:16 +0100240 continue
Tim Hallb9b515c2020-11-01 21:27:19 +0000241 rng = lr_graph.get_or_create_range(tens, cpu_tensor_alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100242 rng.mark_usage(time_for_pass)
243
Fredrik Svedbergf3c7d552022-11-04 09:48:49 +0100244 op = cps.passes[0].ops[0] if cps.passes[0].ops else None
245 op_subgraph = op.attrs.get("subgraph", None) if op else None
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200246
Johan Alfvén673683b2022-09-05 09:39:47 +0200247 if op_subgraph is not None and MemType.Permanent_CPU not in target_mem_type_set:
Fredrik Svedbergf3c7d552022-11-04 09:48:49 +0100248 if op.type == Op.CustomNpuOp:
Johan Alfvén673683b2022-09-05 09:39:47 +0200249 # If the primary-op is an NpuOp that means this is where an Npu subgraph
250 # is called. Go into said subgraph and extract live ranges before continuing.
251 # Use default allocation alignment of 16 for Npu tensors
252 lr_graph = _extract_live_ranges_from_schedule(
253 op_subgraph, target_mem_area, target_mem_type_set, lr_graph
254 )
255 else:
256 # The op has one or more subgraphs in it (a typical op is the While op)
257 # Go into all subgraphs and extract live ranges before continuing.
258 for op_sg in op_subgraph:
259 lr_graph = extract_live_ranges_from_cascaded_passes(
260 op_sg, target_mem_area, target_mem_type_set, lr_graph, cpu_tensor_alignment
261 )
Tim Hall79d07d22020-04-27 18:20:16 +0100262 # Set the new time after handling the Npu subgraph
263 time_for_pass = lr_graph.current_time
264 cps.time = time_for_pass
265
Patrik Gustavssona151f592020-10-16 13:59:52 +0200266 for tens in cps.intermediates + cps.outputs:
Fredrik Svedberg0ae28482021-10-27 13:58:03 +0200267 if tensor_should_be_ignored(tens, target_mem_area, target_mem_type_set):
Tim Hall79d07d22020-04-27 18:20:16 +0100268 continue
Tim Hallb9b515c2020-11-01 21:27:19 +0000269 rng = lr_graph.get_or_create_range(tens, cpu_tensor_alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100270 rng.mark_usage(time_for_pass)
271
Tim Hall79d07d22020-04-27 18:20:16 +0100272 lr_graph.current_time += 2
273
274 end_time = 0
275 for rng in lr_graph.ranges.values():
276 # Find the maximum end time of all live-ranges in the graph
277 end_time = max(end_time, rng.end_time)
278
279 for tens in sg.output_tensors:
Fredrik Svedberg0ae28482021-10-27 13:58:03 +0200280 if tensor_should_be_ignored(tens, target_mem_area, target_mem_type_set):
Tim Hall79d07d22020-04-27 18:20:16 +0100281 continue
Tim Hallb9b515c2020-11-01 21:27:19 +0000282 rng = lr_graph.get_or_create_range(tens, cpu_tensor_alignment)
Tim Hall79d07d22020-04-27 18:20:16 +0100283 rng.mark_usage(end_time)
284
Fredrik Svedberg33c01e62023-02-13 11:32:12 +0100285 # Variable tensor live-range is for entire inference
286 for tens, rng in lr_graph.ranges.items():
287 if tens.is_variable:
288 rng.mark_usage(0, end_time + 1)
289
Tim Hall79d07d22020-04-27 18:20:16 +0100290 # Add subgraph to set of processed subgraphs
291 lr_graph.processed_subgraphs.add(sg)
292 return lr_graph
Tim Halld8339a72021-05-27 18:49:40 +0100293
294
295def create_linear_live_range_graph(sg, target_mem_area, target_mem_type_set, lr_graph):
296 assert lr_graph is not None
297 sg_time = lr_graph.current_time
298 for ps in sg.passes:
299 for tens in ps.inputs + ps.outputs + ps.intermediates:
300 if tens.purpose == TensorPurpose.Weights or tensor_should_be_ignored(
Fredrik Svedberg0ae28482021-10-27 13:58:03 +0200301 tens, target_mem_area, target_mem_type_set
Tim Halld8339a72021-05-27 18:49:40 +0100302 ):
303 continue
Tim Halld8339a72021-05-27 18:49:40 +0100304 rng = lr_graph.get_or_create_range(tens)
305 rng.mark_usage(sg_time)
306
Jacob Bohlin98bfecd2021-06-21 17:22:20 +0200307 for _, op_info in sg.schedule.cost_map.items():
Tim Halld784af72021-06-08 21:25:57 +0100308 for tensor in [op_info.npu_weights_tensor, op_info.npu_scales_tensor]:
Fredrik Svedberg0ae28482021-10-27 13:58:03 +0200309 if tensor and not (tensor_should_be_ignored(tensor, target_mem_area, target_mem_type_set)):
Tim Halld784af72021-06-08 21:25:57 +0100310 rng = lr_graph.get_or_create_range(tensor)
311 rng.mark_usage(sg_time)
Tim Halld8339a72021-05-27 18:49:40 +0100312
313 lr_graph.current_time += 1
314 return lr_graph
315
316
317def _extract_live_ranges_from_schedule(sg, target_mem_area, target_mem_type_set, lr_graph):
318 time_for_cascade = {}
319 for sched_op in sg.sched_ops:
320 op_info = sg.schedule.cost_map[sched_op]
321 cascade = op_info.cascade
322 cascade_info = sg.schedule.cascades.get(cascade, None)
323
Johan Alfvén783d3642022-07-19 14:03:27 +0200324 if cascade_info is None:
325 # Op is not part of a cascade, check if the ifm can be overwritten by the ofm
326 merge_elementwise_op_ranges(sg, sched_op, lr_graph, target_mem_area, target_mem_type_set)
327
Tim Halld8339a72021-05-27 18:49:40 +0100328 time_to_set = time_for_cascade.get(cascade, lr_graph.current_time)
329
330 op_info.time_index = time_to_set
331
332 # Mark usage for all tensors related to this Pass
333 ps = sched_op.parent_ps
334 for tens in ps.inputs + ps.outputs + ps.intermediates:
335 if (
336 target_mem_area == MemArea.Sram
337 and cascade_info
338 and tens == ps.ifm_tensor
339 and sched_op in cascade_info.buffers
340 ):
341 # This tensor is a rolling buffer in a cascade and the size of the LiveRange needs to be modified
342 # for enabling temporal memory snapshots without modifying the original Tensor
343 rng = lr_graph.get_or_create_range(tens)
344 rng.set_buffer_size(cascade_info.buffers[sched_op].elements() * sched_op.ifm.dtype.size_in_bytes())
345 elif (
346 tens.purpose == TensorPurpose.Weights
347 or tens.purpose == TensorPurpose.FSBias
348 or tens.mem_type not in target_mem_type_set
349 or tens.mem_area != target_mem_area
350 ):
351 continue
352
353 else:
354 rng = lr_graph.get_or_create_range(tens)
355
356 rng.mark_usage(time_to_set)
357
Rickard Bolinfd8b5002022-05-16 09:11:06 +0000358 for idx, weight_tens in enumerate(op_info.buffered_weight_tensors):
359 if weight_tens.mem_type in target_mem_type_set and weight_tens.mem_area == target_mem_area:
360 rng = lr_graph.get_or_create_range(weight_tens)
361 start_time = time_to_set
362 length = 1
363 if weight_tens.pre_buffer:
364 start_time -= 1
365 length += 1
366 if len(op_info.buffered_weight_tensors) > 1:
367 last_idx = len(op_info.ofm_depth_slices) % len(op_info.buffered_weight_tensors)
368 # Double buffering: reduce end time of the buffer that is not used last
369 if last_idx != idx:
370 length -= 1
371 rng.mark_usage(start_time, length)
Tim Halld8339a72021-05-27 18:49:40 +0100372
373 if time_to_set == lr_graph.current_time:
374 lr_graph.current_time += 2
375
376 if cascade != 0:
377 time_for_cascade[cascade] = time_to_set
378
379 end_time = lr_graph.update_endtime()
380
381 for tens in sg.output_tensors:
382 if tens.mem_type not in target_mem_type_set or tens.mem_area != target_mem_area:
383 continue
384 rng = lr_graph.get_or_create_range(tens)
385 rng.mark_usage(end_time)
386
387 return lr_graph