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.
In general, the following steps have to performed:
Currently, two use cases are supported:
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.
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
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:
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:
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:
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.