pytorch

1cec7194 - [fx][graph opts] port FoldLayerNormArithmetic from glow to FX (#69715)

[fx][graph opts] port FoldLayerNormArithmetic from glow to FX (#69715) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/69715 This optimization folds simple arithmetic operations following a layer norm operator into the layer norm bias and variance. Layer norm is `output = stuff * weight + bias.` if we add/subtract multiply/divide the output, we can optimize this into weight and bias. With constant fold (another graph opt), this eliminates the arithmetic op from the computational graph. e.g. `x - output = x - (stuff * weight + bias) = stuff * (-weight) + (x - bias)` Test Plan: Added unit tests ``` buck test mode/opt glow/fb/fx/graph_opts:test_fx_graph_opts -- test_fold_layer_norm_arithmetic BUILD SUCCEEDED Tpx test run coordinator for Facebook. See https://fburl.com/tpx for details. Running with tpx session id: 25eb3f06-2976-471b-ac07-eea3e42ecf7c Trace available for this run at /tmp/tpx-20211207-152940.591324/trace.log RemoteExecution session id: reSessionID-25eb3f06-2976-471b-ac07-eea3e42ecf7c-tpx Started reporting to test run: https://www.internalfb.com/intern/testinfra/testrun/3940649754647173 ✓ ListingSuccess: glow/fb/fx/graph_opts:test_fx_graph_opts - main (4.238) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_00_identity (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.016) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_03_Add_3_5_LayerNorm_3_5_ (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.102) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_10__LayerNorm (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.254) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_04_Add_1_3_5_LayerNorm_3_5_ (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.230) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_07_Div_LayerNorm (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.316) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_08_unoptimizable_Div_LayerNorm (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.345) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_05_Add_4_3_5_LayerNorm_3_5_ (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.250) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_06_Mul_LayerNorm (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.286) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_09_Mul_3_5_LayerNorm (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.122) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_02_Sub_LayerNorm_3_5_ (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.262) ✓ Pass: glow/fb/fx/graph_opts:test_fx_graph_opts - test_fold_layer_norm_arithmetic_01_Add_LayerNorm_3_5_ (glow.fb.fx.graph_opts.tests.test_fx_graph_opts.TestFoldLayerNormArithmetic) (0.231) Summary Pass: 11 ListingSuccess: 1 If you need help understanding your runs, please follow the wiki: https://fburl.com/posting_in_tpx_users Finished test run: https://www.internalfb.com/intern/testinfra/testrun/3940649754647173 ``` Reviewed By: jfix71 Differential Revision: D32778848 fbshipit-source-id: ab8b3dd7465c1cf081041f3589ad107ffc90b9df (cherry picked from commit 5fb72184f4384280104a1ebc3600d7d7cd29e337)