Make AIT use fp32 accumulation for reduce_3d kernels by default. (fac…

…ebookincubator#862)

Summary:
Pull Request resolved: facebookincubator#862

Use fp32 accumulation by default, only use fp16 if use_fp16_acc == True.

Differential Revision: D47928197

fbshipit-source-id: 7d321cbc2371b4bda3011c659690c09002a7f40f

python/aitemplate/backend/cuda/reduce/reduce_3d.py

@@ -28,6 +28,7 @@
 from aitemplate.backend.common import tensor_accessor_codegen
 
 from aitemplate.backend.cuda.reduce import reduce_small_axis
+from aitemplate.backend.target import Target
 
 
 DEFAULT_PROLOGUE_TEMPLATE = jinja2.Template(
@@ -830,7 +831,13 @@ def gen_function(
     output_type = backend_spec.dtype_to_lib_type(y._attrs["dtype"])
     if accumulation_type is None:
         # follow pytorch's semantics
-        acc_type = output_type
+        if (
+            Target.current()._kwargs.get("use_fp16_acc", False)
+            and y._attrs["dtype"] == "float16"
+        ):
+            acc_type = output_type
+        else:
+            acc_type = "float"
     else:
         acc_type = accumulation_type
 

python/aitemplate/backend/cuda/reduce/var.py

@@ -24,6 +24,7 @@
 from aitemplate.backend import registry
 from aitemplate.backend.backend_spec import CUDASpec
 from aitemplate.backend.cuda.reduce import reduce_3d
+from aitemplate.backend.target import Target
 
 
 EXTRA_CODE_TEMPLATE = jinja2.Template(
@@ -294,10 +295,12 @@ def var_gen_function(func_attrs) -> str:
     """
     bessel = "true" if func_attrs["unbiased"] else "false"
     backend_spec = CUDASpec()
-    elem_output_type = backend_spec.dtype_to_lib_type(
-        func_attrs["outputs"][0]._attrs["dtype"]
-    )
-    acc_type = f"WelfordData<{elem_output_type}, {bessel}>"
+    output_type = func_attrs["outputs"][0]._attrs["dtype"]
+    elem_output_type = backend_spec.dtype_to_lib_type(output_type)
+    if Target.current()._kwargs.get("use_fp16_acc", False) and output_type == "float16":
+        acc_type = f"WelfordData<{elem_output_type}, {bessel}>"
+    else:
+        acc_type = f"WelfordData<float, {bessel}>"
     return reduce_3d.gen_function(
         func_attrs,
         "cutlass::welford_op",