1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
|
# AutoHeuristic
AutoHeuristic is a framework that allows one to use results from autotuning to learn a heuristic as a decision tree, that can be generated to code and shipped with compiler.
## How to use AutoHeuristic
In general, the following steps have to performed:
- The AutoHeursitic constructor has to be called.
- A script that runs benchmarks in order to collect training data has to be implemented.
- The train_decision.py (if you want to learn a decision tree) or train_regression.py (if you want to learn a regression tree) script has to be run in order to learn the heuristic and generate it to code.
## Step 1: Calling the AutoHeuristic constructor
Currently, two use cases are supported:
### Use case 1: Local autotuning
When your feedback function is able to immediately return a result, you can just call the AutoHeuristic constructor. This is done e.g. for pad_mm
```
autoheuristic = AutoHeuristic(
fallback=fallback,
choices=choices,
feedback=feedback,
context=context,
name=name,
augment_context=pad_mm_operations(),
precondition=pad_mm_precondition,
)
```
Here, `feedback` is a function that benchmarks a given choice and returns the execution time. For an example, see: https://github.com/pytorch/pytorch/blob/main/torch/_inductor/fx_passes/pad_mm.py.
### Use case 2: Kernel choice selection
If you want to use AutoHeuristic for kernel choice selection, you have to call the AutoHeuristicSelectAlgorithm constructor. This is done e.g. for mixed_mm
```
autoheuristic = AutoHeuristicSelectAlgorithm(
fallback=fallback,
choices=choices,
input_nodes=input_nodes,
context=context,
name=name,
augment_context=ops,
precondition=precondition,
)
```
This call has to be followed by a call to `autotune_select_algorithm()`,
```
autotune_select_algorithm(name, choices, input_nodes, layout)
```
Note that `choices`, `input_nodes`, and `name` in the `AutoHeuristicSelectAlgorithm()` and `autotune_select_algorithm()` calls have to match when you want to use AutoHeuristic to collect data.
For an example, see: https://github.com/pytorch/pytorch/blob/main/torch/_inductor/kernel/mm.py
## Step 2: Collecting training data
After adding the call to the AutoHeuristic constructor, you need to collect training data in order to learn a heuristic. Let's say you have a script `run.py` that triggers the AutoHeuristic constructor that you just added. Run the following command in order to store data into file `train.txt`:
```
TORCHINDUCTOR_AUTOHEURISTIC_LOG_PATH="train.txt" \
TORCHINDUCTOR_AUTOHEURISTIC_COLLECT="pad_mm" python run.py
```
Replace "pad_mm" with the name you provided in the call to the AutoHeuristic constructor.
AutoHeuristic provides a `BenchmarkRunner` class (https://github.com/pytorch/pytorch/blob/main/torchgen/_autoheuristic/benchmark_runner.py) that simplifies the process of collecting data. To use it, create a new class that subclasses `BenchmarkRunner`, and implements the `run_benchmark()` and `create_input()` methods.
These examples might be helpful:
- https://github.com/pytorch/pytorch/blob/main/torchgen/_autoheuristic/pad_mm/gen_data_pad_mm.py
- https://github.com/pytorch/pytorch/blob/main/torchgen/_autoheuristic/mixed_mm/gen_data_mixed_mm.py
## Step 3: Learning a heuristic and using it
Once you have collected enough training data, you are ready to learn a heuristic:
```
python torchgen/_autoheuristic/train_decision.py train.txt --heuristic-name SimpleHeuristic
```
will learn a heuristic and generate it to `torch/_inductor/autoheuristic/artifacts/_SimpleHeuristic.py`.
You can now use your learned heuristic:
```
TORCHINDUCTOR_AUTOHEURISTIC_USE="pad_mm" python run.py
```
Here, you again have to replace "pad_mm" with the name you provided in the call to the AutoHeuristic constructor.
Instead of just running the `train_decision.py` script, you probably want to customize the training process in some way. To do this, create a new class that subclasses `AHTrainDecision` and override methods you want to customize. Here are some examples:
- https://github.com/pytorch/pytorch/blob/main/torchgen/_autoheuristic/mixed_mm/train_decision_mixedmm.py
- https://github.com/pytorch/pytorch/blob/main/torchgen/_autoheuristic/pad_mm/train_decision_pad_mm.py
## Other
### How do I specify features that the heuristic is going to use to make a decision?
The AutoHeuristic constructor requires a `context` argument of type `AHContext`, which will contain all features. You specify features in the following way:
```
context = AHContext()
# adding numerical features
context.add_feature("m", mat1.shape[0])
context.add_feature("k", mat1.shape[1])
# adding a categorical feture
context.add_feature("mat1_dtype", mat1.dtype, is_categorical=True)
```
You might want to use features that are a combination of other features, such as `m*k`. You can of course add such features in the same way as above, i.e.,
```
context.add_feature("m*k", mat1.shape[0] * mat1.shape[1])
```
but AutoHeuristic also provides a way to 'augment' features. Augmented features are not stored when data is collected, instead they are created before a heuristic is learned, or before a learned heuristic is used. You can specify such augmented features by creating a list of `AHOperation` objects:
```
def m_times_k(data: Any) -> float:
return data['m'] * data['k']
m_times_k_op = AHOperation("m*k', m_times_k)
ah_operations = [m_times_k_op]
# specify augmented features by setting `augment_context` to `ah_operations`
autoheuristic = AutoHeuristic(..., augment_context=ah_operations, ...)
```
Note that you also have to specify these operations when you want to learn a heuristic. Look at the `add_new_features()` method in these examples, to see how it is done:
- https://github.com/pytorch/pytorch/blob/main/torchgen/_autoheuristic/mixed_mm/train_decision_mixedmm.py
- https://github.com/pytorch/pytorch/blob/main/torchgen/_autoheuristic/pad_mm/train_decision_pad_mm.py
### Where has AutoHeuristic already been used?
Take a look at the following PRs in which AutoHeuristic has enabled for various optimizations.
Looking at these examples may be helpful if you want to use AutoHeuristic yourself.
- pad_mm: https://github.com/pytorch/pytorch/pull/128643
- mixed_mm:
- Enabling of AutoHeuristic: https://github.com/pytorch/pytorch/pull/131610
- Script to collect data: https://github.com/pytorch/pytorch/pull/131611
- A100 heuristic: https://github.com/pytorch/pytorch/pull/131613
- H100 heuristic: https://github.com/pytorch/pytorch/pull/132685
- flex_attention: https://github.com/pytorch/pytorch/pull/130398
- mm (heuristic for ranking choices):
- https://github.com/pytorch/pytorch/pull/131615
- https://github.com/pytorch/pytorch/pull/131617
- https://github.com/pytorch/pytorch/pull/131705
- https://github.com/pytorch/pytorch/pull/131714
|
