Adding a flag to enable CPU fusion in benchmarks (pytorch#48612)

Summary: Pull Request resolved: pytorch#48612

Test Plan: python -m benchmarks.tensorexpr --device cpu --mode fwd --jit_mode trace --cpu_fusion element

Reviewed By: heitorschueroff

Differential Revision: D26548643

Pulled By: navahgar

fbshipit-source-id: adb537818d77c9b6b0fe434ae6d963a5f348ad24

Author: navahgar

benchmarks/tensorexpr/__main__.py

@@ -116,6 +116,12 @@ def main():
         action='store_true',
         help="Disable shape randomization in dynamic benchmarks.",
     )
+    parser.add_argument(
+        "--cpu_fusion",
+        default=False,
+        action='store_true',
+        help="Enable CPU fusion.",
+    )
 
     args = parser.parse_args()
 
@@ -139,6 +145,13 @@ def main():
     else :
         raise ValueError("Undefined fuser: {}".format(args.cuda_fuser))
 
+    if args.cpu_fusion:
+        import torch
+        torch._C._jit_override_can_fuse_on_cpu(True)
+    else:
+        import torch
+        torch._C._jit_override_can_fuse_on_cpu(False)
+
     def set_global_threads(num_threads):
         os.environ["OMP_NUM_THREADS"] = str(num_threads)
         os.environ["MKL_NUM_THREADS"] = str(num_threads)