Skip to content

Commit 6048566

Browse files
committed
update memory planning docs
1 parent 329184a commit 6048566

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

docs/source/compiler-memory-planning.md

+6-6
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ MemoryPlanning is the very last action taken before taking an `ExportedProgram`
99
Concretely, there are three passes related to memory planning:
1010
* `SpecPropPass` computes a TensorSpec for each tensor in the graph (inputs, intermediates or outputs). The most important field of the tensor spec is a symbolic expression of the shapes of the tensor, where the initial set of symbols comes from the dimensions of input tensors, intermediate tensor shapes’ symbolic expression is propagated via tensor operations. The dimensions can be marked as either dynamic or static by users and when the dims are dynamic, users are required to annotate the dim with a ValueRange.
1111

12-
* `SymShapEvalPass` evaluates the symbolic expressions to concrete integers with their upper bounds. There are two ways to doing the upper bound specialization:
12+
* `SymShapeEvalPass` evaluates the symbolic expressions to concrete integers with their upper bounds. There are two ways to doing the upper bound specialization:
1313
HintBasedSymShapeEval (to be deprecated) is the old way of evaluating the upper bound. It doesn’t look at the ValueRange of the symbols but uses the shapes of example inputs to replace all the symbols. We call it “hint based“ because the example inputs’ shapes are just hints of what the input shapes might be at run time and are used for tracing only. ValueRangeBasedSymShapeEval is the recommended way of doing UpperBoundMemory planning. It will actually look at the ValueRange of the symbols and do an inference over the ranges to get a real upper bound.
1414

1515
* `MemoryPlanningPass` does the actual memory planning given all tensors get a TensorSpec with concrete integer shapes.
@@ -18,9 +18,9 @@ HintBasedSymShapeEval (to be deprecated) is the old way of evaluating the upper
1818

1919
ExecuTorch provides two options for memory planning algorithms out of the box, but users can define their own if the provided options are inappropriate or insufficient for their use case.
2020

21-
* The naive algorithm simply concatenates all the tensors together in a linear memory without considering any memory re-use. It serves as an upper bound for total memory consumption and serves as a baseline.
21+
* The naive algorithm simply concatenates all the tensors together in a linear memory block without considering memory re-use. It serves as an upper bound for total memory consumption and serves as a baseline.
2222

23-
* The Greedy algorithm tries to re-use the already allocated memory and choose based on the best-fit criteria. Specifically:
23+
* The Greedy algorithm tries to re-use the already allocated memory based on the best-fit criteria. Specifically:
2424
When there isn’t an allocated memory whose lifetime doesn’t overlap with the current tensor that we try to do memory planning for, we allocate a new memory buffer with the same size and lifetime as the current tensor. When there is one or more allocated memory buffer, whose lifetime overlaps with the current tensor, we pick the buffer that has the closest size with current tensor so as to reduce memory fragmentation. Finally, we allocate these memory buffers linearly in memory.
2525

2626

@@ -48,7 +48,7 @@ Users can write custom memory plans to take advantage of multiple memory locatio
4848

4949
```python
5050
class CustomPoolMemoryPlanningPass(MemoryPlanningPass):
51-
def call(self, graph_module: GraphModule) -> PassResult:
51+
def run(self, graph_module: GraphModule, graph_signature: Optional[ExportGraphSignature]) -> PassResult:
5252
for subgm in graph_module.modules():
5353
if not isinstance(subgm, GraphModule):
5454
continue
@@ -68,7 +68,7 @@ class CustomPoolMemoryPlanningPass(MemoryPlanningPass):
6868
elif node.target == torch.ops.aten.mul.out:
6969
node.meta["spec"].mem_id = 1
7070

71-
return super().call(graph_module)
71+
return super().run(graph_module, graph_signature)
7272
```
7373

7474
Then later when lowering to ExecuTorch you can use your custom plan in the following way:
@@ -83,4 +83,4 @@ program = edge_program.to_executorch(
8383
)
8484
```
8585

86-
Users attempting to write a custom memory planning algorithm should start by looking at [the greedy algorithm's implementation](https://github.com/pytorch/executorch/blob/d62c41ca86435e5316e7ed292b6d68aff27a2fb7/exir/memory_planning.py#L459C1-L459C12)
86+
Users attempting to write a custom memory planning algorithm should start by looking at [the greedy algorithm's implementation](https://github.com/pytorch/executorch/blob/d62c41ca86435e5316e7ed292b6d68aff27a2fb7/exir/memory_planning.py#L459C1-L459C12).

0 commit comments

Comments
 (0)