1#pragma once
2
3#include <torch/csrc/jit/ir/ir.h>
4
5namespace torch {
6namespace 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.
67TORCH_API bool PeepholeOptimizeListIdioms(
68 const std::shared_ptr<Graph>& graph,
69 bool refine_list_len = false);
70
71} // namespace jit
72} // namespace torch
73