Introduce AIT_USE_FAST_MATH environment flag (facebookincubator#627)

Summary:
Pull Request resolved: facebookincubator#627

Whether the fast math option should be used for the device code generation. Fast math implies the use of approximate math operations (say, a division operation), allowing to gain speed at the cost of accuracy. Default value is "1".

Reviewed By: aakhundov

Differential Revision: D45321954

fbshipit-source-id: 9df3583eef5f7284176d338100f6702bace07a90

docs/source/reference/env.rst

@@ -48,3 +48,5 @@ Miscellaneous
 **LOGLEVEL**: It is used to control the logging level in Python. The default value is "INFO". "DEBUG" is useful for debugging.
 
 **AIT_PLOT_SHORTEN_TENSOR_NAMES**: If set to "1", shorten too long tensor names for a plot of a model graph, thus making a plot much easier to analyze visually. "0" by default.
+
+**AIT_USE_FAST_MATH**: If set to "0", no fast math option will be used for the device code generation. Default value is "1".

python/aitemplate/backend/cuda/target_def.py

@@ -120,11 +120,12 @@ def _build_compile_options(self):
             environ.get_compiler_opt_level(),
             "-std=c++17",
             "--expt-relaxed-constexpr",
-            "--use_fast_math",
             f"-I{ait_static_path}",
         ] + ["-I" + path for path in cutlass_path]
         if self._ndebug == 1:
             options.append("-DNDEBUG")
+        if environ.use_fast_math():
+            options.append("--use_fast_math")
         return " ".join(options)
 
     def src_extension(self):
@@ -277,7 +278,6 @@ def _build_compile_options(self):
                     "-DCUTLASS_USE_TANH_FOR_SIGMOID=1",
                     "-w",
                     "--expt-relaxed-constexpr",
-                    "--use_fast_math",
                     f"-gencode=arch=compute_{nvcc_arch},code=[sm_{nvcc_arch},compute_{nvcc_arch}]",
                     "-Xcompiler=-Wconversion",
                     environ.get_compiler_opt_level(),
@@ -286,6 +286,8 @@ def _build_compile_options(self):
             )
             if self._ndebug == 1:
                 options.append("-DNDEBUG")
+            if environ.use_fast_math():
+                options.append("--use_fast_math")
             FBCUDA.compile_options_ = " ".join(options)
         compile_options = FBCUDA.compile_options_
         _LOGGER.info(f"The compile options are: {compile_options}")

python/aitemplate/utils/environ.py

@@ -36,6 +36,16 @@ def get_compiler_opt_level() -> str:
     return compiler_opt
 
 
+def use_fast_math() -> str:
+    """
+    Whether the fast math option should be used for the device code generation.
+    Fast math implies the use of approximate math operations (say,
+    a division operation), allowing to gain speed at the cost of accuracy.
+    Default value is "1".
+    """
+    return os.getenv("AIT_USE_FAST_MATH", "1") == "1"
+
+
 def force_profiler_cache() -> bool:
     """
     Force the profiler to use the cached results. The profiler will throw