Repository: pytorch/TensorRT

Issue:
✨[Feature] Add API to disable TF32 in TRTorch 
There is a builder flag to disable TF32 format in TensorRT. Disabling it will ensure tactic optimizer picks FP32 format.


Discussion:
What is the TensorRT API to do this? 
Here is the [flag](https://gitlab-master.nvidia.com/TensorRT/TensorRT/-/blob/rel-7.2/caffeHarness/infer_caffe/caffeHarness.cpp#L1306) corresponding to trtexec code  which has `noTF32` flag
```
config->clearFlag(BuilderFlag::kTF32);
```

This issue has not seen activity for 30 days, Remove stale label or comment or this will be closed in 5 days

Resolving patch:
diff --git a/py/trtorch/_compile_spec.py b/py/trtorch/_compile_spec.py
--- a/py/trtorch/_compile_spec.py
+++ b/py/trtorch/_compile_spec.py
@@ -135,6 +135,10 @@ def _parse_compile_spec(compile_spec: Dict[str, Any]) -> trtorch._C.CompileSpec:
     if "op_precision" in compile_spec:
         info.op_precision = _parse_op_precision(compile_spec["op_precision"])
 
+    if "disable_tf32" in compile_spec:
+        assert isinstance(compile_spec["disable_tf32"], bool)
+        info.disable_tf32 = compile_spec["disable_tf32"]
+
     if "refit" in compile_spec:
         assert isinstance(compile_spec["refit"], bool)
         info.refit = compile_spec["refit"]
@@ -201,6 +205,7 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
                             "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
                         },
                         "op_precision": torch.half, # Operating precision set to FP16
+                        "disable_tf32": False, # Force FP32 layers to use traditional as FP32 format vs the default behavior of rounding the inputs to 10-bit mantissas before multiplying, but accumulates the sum using 23-bit mantissas
                         "refit": False, # enable refit
                         "debug": False, # enable debuggable engine
                         "strict_types": False, # kernels should strictly run in operating precision
@@ -239,6 +244,7 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
 
     backend_spec.set_device(d)
     backend_spec.set_op_precision(int(parsed_spec.op_precision))
+    backend_spec.set_disable_tf32(parsed_spec.disable_tf32)
     backend_spec.set_refit(parsed_spec.refit)
     backend_spec.set_debug(parsed_spec.debug)
     backend_spec.set_refit(parsed_spec.refit)
diff --git a/py/trtorch/_compiler.py b/py/trtorch/_compiler.py
--- a/py/trtorch/_compiler.py
+++ b/py/trtorch/_compiler.py
@@ -99,6 +99,7 @@ def convert_method_to_trt_engine(module: torch.jit.ScriptModule, method_name: st
                         "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
                     },
                     "op_precision": torch.half, # Operating precision set to FP16
+                    "disable_tf32": False, # Force FP32 layers to use traditional as FP32 format vs the default behavior of rounding the inputs to 10-bit mantissas before multiplying, but accumulates the sum using 23-bit mantissas
                     "refit": false, # enable refit
                     "debug": false, # enable debuggable engine
                     "strict_types": false, # kernels should strictly run in operating precision


Test patch:
diff --git a/tests/py/test_api.py b/tests/py/test_api.py
--- a/tests/py/test_api.py
+++ b/tests/py/test_api.py
@@ -20,7 +20,8 @@ def test_compile_traced(self):
                 "device_type": trtorch.DeviceType.GPU,
                 "gpu_id": 0,
                 "dla_core": 0,
-                "allow_gpu_fallback": False
+                "allow_gpu_fallback": False,
+                "disable_tf32": False
             }
         }
 
@@ -35,7 +36,8 @@ def test_compile_script(self):
                 "device_type": trtorch.DeviceType.GPU,
                 "gpu_id": 0,
                 "dla_core": 0,
-                "allow_gpu_fallback": False
+                "allow_gpu_fallback": False,
+                "disable_tf32": False
             }
         }
 
diff --git a/tests/py/test_to_backend_api.py b/tests/py/test_to_backend_api.py
--- a/tests/py/test_to_backend_api.py
+++ b/tests/py/test_to_backend_api.py
@@ -29,6 +29,7 @@ def setUp(self):
                     "num_min_timing_iters": 2,
                     "num_avg_timing_iters": 1,
                     "max_batch_size": 0,
+                    "disable_tf32": False,
                 })
         }