1#include <c10/util/irange.h>
2#include <torch/csrc/jit/ir/alias_analysis.h>
3#include <torch/csrc/jit/ir/ir_views.h>
4#include <torch/csrc/jit/passes/frozen_concat_linear.h>
5#include <torch/csrc/jit/passes/frozen_conv_folding.h>
6#include <torch/csrc/jit/passes/frozen_graph_optimizations.h>
7#include <torch/csrc/jit/passes/frozen_linear_folding.h>
8#include <torch/csrc/jit/passes/remove_dropout.h>
9#include <torch/csrc/jit/runtime/graph_executor.h>
10#include <torch/csrc/utils/memory.h>
11
12namespace torch {
13namespace jit {
14
15void OptimizeFrozenGraph(
16 std::shared_ptr<Graph>& graph,
17 bool optimize_numerics) {
18 removeDropout(graph);
19 FrozenConcatLinear(graph);
20 // run a couple times to capture Conv -> Mul -> Add etc
21 if (optimize_numerics) {
22 bool changed = false;
23 do {
24 changed = false;
25 changed |= FoldFrozenConvBatchnorm(graph);
26 changed |= FoldFrozenConvAddOrSub(graph);
27 changed |= FoldFrozenConvMulOrDiv(graph);
28 changed |= FoldFrozenLinearBatchnorm(graph);
29 } while (changed);
30 }
31}
32
33} // namespace jit
34} // namespace torch
35