[trt] Algorithm recorder/replayer (#4)

Summary:
Pull Request resolved: pytorch/pytorch-canary#4

Pull Request resolved: pytorch#67211

Record the algorithm selection, dump it in json format and replay it. This has potential to
1. consistently repro the issue (algo selection could be sensitive to local benchmark timing)
2. manual edit the dumped json file to control algorithm selection.

Reviewed By: wushirong, 842974287

Differential Revision: D31888836

fbshipit-source-id: 4611fda548f7391776f1ad61572b1f59fa4665b6

Author: Yinghai Lu

torch/fx/experimental/fx2trt/fx2trt.py

@@ -414,6 +414,7 @@ def run(
         int8_mode=False,
         force_fp32_output=False,
         strict_type_constraints=False,
+        algorithm_selector=None,
     ) -> TRTInterpreterResult:
         # For float outputs, we set their dtype to fp16 only if fp16_mode=True and
         # force_fp32_output=False.
@@ -444,6 +445,10 @@ def run(
             for optimization_profile in self.optimization_profiles:
                 builder_config.add_optimization_profile(optimization_profile)
 
+        if algorithm_selector:
+            builder_config.set_flag(trt.BuilderFlag.DISABLE_TIMING_CACHE)
+            builder_config.algorithm_selector = algorithm_selector
+
         engine = self.builder.build_engine(self.network, builder_config)
         assert engine
         return engine, self._input_names, self._output_names