Skip to content

Commit ce51fe5

Browse files
sabarishsnkfacebook-github-bot
authored andcommitted
Update CadenceMemoryPlanning to support per-memory alignment constraint
Summary: Currently `CadenceMemoryPlanning` only supports a single `mem_alignment` parameter that it applies to all memories in the hierarchy. This diff moves the alignment constraint to the `MemoryConfig` definition and makes it per-memory. Reviewed By: dulinriley Differential Revision: D69881079
1 parent 88b3441 commit ce51fe5

File tree

4 files changed

+26
-12
lines changed

4 files changed

+26
-12
lines changed

backends/cadence/aot/compiler.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,6 @@ def export_to_executorch_gen_etrecord(
264264
alloc_graph_output: bool = True,
265265
memory_config: Optional[MemoryConfig] = None,
266266
dump_graphs: bool = False,
267-
mem_alignment: int = 1,
268267
) -> ExecutorchProgramManager:
269268
cadence_passes = get_cadence_passes(opt_level)
270269
edge_prog_manager = export_to_edge(model, inputs, dump_graphs)
@@ -291,7 +290,6 @@ def export_to_executorch_gen_etrecord(
291290
mem_algo=mem_algo,
292291
alloc_graph_input=alloc_graph_input,
293292
alloc_graph_output=alloc_graph_output,
294-
mem_alignment=mem_alignment,
295293
)
296294

297295
# Get executorch program after Cadence specific passes

backends/cadence/aot/memory_planning.py

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,12 @@ def get_size(memory_config: MemoryConfig, exir_id: int) -> int:
4040
return memory_config.memory_sizes[exir_id - 1]
4141

4242

43+
def get_alignment(memory_config: MemoryConfig, exir_id: int) -> int:
44+
# EXIR's spec.mem_id is indexed from 1..N.
45+
assert memory_config.memory_alignments is not None
46+
return memory_config.memory_alignments[exir_id - 1]
47+
48+
4349
def get_aligned_offset(pre_aligned_offset: int, alignment: int) -> int:
4450
return int(math.ceil(pre_aligned_offset / alignment) * alignment)
4551

@@ -84,6 +90,10 @@ def position_based_greedy_with_hierarchy(
8490
]
8591
] = None,
8692
) -> List[int]:
93+
# We do not use the `alignment` parameter and instead use the per-memory alignment
94+
# constraints from `memory_config`.
95+
del alignment
96+
8797
num_memories = get_num_memories(memory_config)
8898
bufsizes = [0] * num_memories
8999
allocated_buffers: List[List[TensorSpec]] = [[] for _ in range(num_memories)]
@@ -103,7 +113,7 @@ def overlap(spec: TensorSpec) -> Optional[TensorSpec]:
103113

104114
def memory_available(spec: TensorSpec) -> bool:
105115
return get_aligned_offset(
106-
spec.mem_offset + spec.allocated_memory, alignment
116+
spec.mem_offset + spec.allocated_memory, get_alignment(memory_config, spec.mem_id)
107117
) <= get_size(memory_config, spec.mem_id)
108118

109119
# Iterate over all the specs in sorted order
@@ -124,7 +134,7 @@ def memory_available(spec: TensorSpec) -> bool:
124134
spec.mem_offset = 0
125135
while memory_available(spec) and (overlapped := overlap(spec)):
126136
spec.mem_offset = get_aligned_offset(
127-
overlapped.mem_offset + overlapped.allocated_memory, alignment
137+
overlapped.mem_offset + overlapped.allocated_memory, get_alignment(memory_config, spec.mem_id)
128138
)
129139
if memory_available(spec):
130140
allocated_buffers[spec.mem_id].append(spec)
@@ -172,6 +182,10 @@ def greedy_by_size_for_offset_calculation_with_hierarchy(
172182
]
173183
] = None,
174184
) -> List[int]:
185+
# We do not use the `alignment` parameter and instead use the per-memory alignment
186+
# constraints from `memory_config`.
187+
del alignment
188+
175189
num_memories = get_num_memories(memory_config)
176190
bufsizes = [0] * num_memories
177191
allocated_buffers = [[] for _ in range(num_memories)]
@@ -213,13 +227,13 @@ def greedy_by_size_for_offset_calculation_with_hierarchy(
213227
prev_offset = max(
214228
get_aligned_offset(
215229
allocated_spec.mem_offset + allocated_spec.allocated_memory,
216-
alignment,
230+
get_alignment(memory_config, spec.mem_id),
217231
),
218232
prev_offset,
219233
)
220234
if spec.mem_offset is None:
221235
if get_aligned_offset(
222-
prev_offset + spec.allocated_memory, alignment
236+
prev_offset + spec.allocated_memory, get_alignment(memory_config, spec.mem_id)
223237
) > get_size(memory_config, spec.mem_id):
224238
continue
225239
else:
@@ -439,7 +453,6 @@ def __init__(
439453
]
440454
]
441455
] = None,
442-
mem_alignment: int = 1,
443456
) -> None:
444457
self._init_mem_algos()
445458

@@ -450,9 +463,6 @@ def __init__(
450463
self.alloc_graph_output = alloc_graph_output
451464
self.additional_constraint_gen_passes = additional_constraint_gen_passes
452465

453-
assert mem_alignment > 0, "mem_alignment must be positive"
454-
self.mem_alignment = mem_alignment
455-
456466
def _init_mem_algos(self) -> None:
457467
self.available_mem_algos = [
458468
position_based_greedy_with_hierarchy,
@@ -489,7 +499,6 @@ def run(
489499
allow_lifetime_and_storage_overlap=(self.opt_level >= 2),
490500
alloc_graph_input=self.alloc_graph_input,
491501
alloc_graph_output=self.alloc_graph_output,
492-
alignment=self.mem_alignment,
493502
)
494503
mem_planning.run(graph_module, graph_signature)
495504

backends/cadence/aot/tests/test_memory_passes.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from executorch.backends.cadence.aot import compiler
1717
from executorch.backends.cadence.aot.memory_planning import find_peak_memory_usage
1818
from executorch.backends.cadence.aot.pass_utils import count_node
19+
from executorch.backends.cadence.aot.utils import MemoryConfig
1920
from executorch.exir import memory
2021
from executorch.exir.dialects._ops import ops as exir_ops
2122
from executorch.exir.memory_planning import collect_specs_from_nodes
@@ -792,7 +793,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
792793
mem_algo=mem_algo,
793794
alloc_graph_input=False,
794795
alloc_graph_output=False,
795-
mem_alignment=37,
796+
memory_config=MemoryConfig(memory_sizes=[0x1000000000], memory_alignments=[37])
796797
)
797798
.exported_program()
798799
.graph_module

backends/cadence/aot/utils.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,13 +256,19 @@ def save_bpte_program(
256256
@dataclass
257257
class MemoryConfig:
258258
memory_sizes: List[int]
259+
# Alignment constraint for each memory region in bytes.
260+
memory_alignments: Optional[List[int]] = None
259261

260262
# Optional fields for logs
261263
memory_names: Optional[List[str]] = None
262264
base_addrs: Optional[List[int]] = None
263265
memory_xml_path: Optional[str] = None
264266
MemorySpace: Optional[enum.Enum] = None
265267

268+
def __post_init__(self) -> None:
269+
if self.memory_alignments is None:
270+
self.memory_alignments = [1] * len(self.memory_sizes)
271+
266272
# get num memories indexed from 1..N, compatible with EXIR's spec.mem_id
267273
def get_num_memories(self) -> int:
268274
return len(self.memory_sizes) + 1

0 commit comments

Comments
 (0)