1 | #pragma once |
2 | |
3 | #include <torch/csrc/jit/ir/ir.h> |
4 | |
5 | namespace torch { |
6 | namespace jit { |
7 | |
8 | // Peephole Optimizes List ops such as len(li) and li[1]. |
9 | // 1. Construct/Unpack optimizations |
10 | // Given a function like this: |
11 | // def foo(a, b): |
12 | // li = [a, b] |
13 | // x, y = li |
14 | // return x, y |
15 | // This pass produces (after dead code elimination): |
16 | // def foo(a, b): |
17 | // return a, b |
18 | // |
19 | // This is only applied to lists that are not modified. |
20 | // |
21 | // 2. getitem optimizations |
22 | // Given a function like this: |
23 | // def foo(a, b): |
24 | // li = [a, b] |
25 | // x = li[0] |
26 | // return x |
27 | // This pass produces (after dead code elimination): |
28 | // def foo(a, b): |
29 | // return a |
30 | // |
31 | // This optimization can only happen if the list is not modified. |
32 | // |
33 | // 3. len optimizations |
34 | // Given a function like this: |
35 | // def foo(): |
36 | // li = [1, 2] |
37 | // return len(li) |
38 | // This pass produces (after dead code elimination): |
39 | // def foo(): |
40 | // return 2 |
41 | // |
42 | // This has the same requirements as the getitem optimizations. |
43 | // |
44 | // 4. ListConstruct + ListConstruct |
45 | // Given a function like this: |
46 | // def foo(): |
47 | // return [1, 2] + [3, 4] |
48 | // This pass produces (after dead code elimination): |
49 | // def foo(): |
50 | // return [1, 2, 3, 4] |
51 | // |
52 | // This is only applied to lists that are not modified. |
53 | // |
54 | // 5. Slice |
55 | // Given a function like this: |
56 | // def foo(): |
57 | // return [1, 2, 3, 4, 5][0:2] |
58 | // This pass produces (after deadcode elimination): |
59 | // def foo(): |
60 | // return [1, 2] |
61 | // |
62 | // Currently this is invoked as part of PeepholeOptimize |
63 | // return true if graph is modified. |
64 | // If `refine_list_len` is true will attempt to refine the len of lists through |
65 | // len comparisons and assertions. This does not generally optimize pytorch |
66 | // programs so it is not called by default in PeepholeOptimize. |
67 | TORCH_API bool PeepholeOptimizeListIdioms( |
68 | const std::shared_ptr<Graph>& graph, |
69 | bool refine_list_len = false); |
70 | |
71 | } // namespace jit |
72 | } // namespace torch |
73 | |