Propagate NaN for layernorm

c466ccd7 · rocking · 798670d8 · c466ccd7
Commit c466ccd7 authored Jul 11, 2022 by rocking
Show whitespace changes
Inline Side-by-side

Showing with 12 additions and 16 deletions

include/ck/tensor_operation/gpu/grid/gridwise_layernorm.hpp include/ck/tensor_operation/gpu/grid/gridwise_layernorm.hpp +12 -16

No files found.
--- a/include/ck/tensor_operation/gpu/grid/gridwise_layernorm.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_layernorm.hpp
@@ -97,22 +97,18 @@ struct GridwiseLayernorm_mk_to_mk
    using ThreadReduceDstDesc_M =
        decltype(make_naive_tensor_descriptor_packed(make_tuple(Number<MThreadSliceSize>{})));

-    using BlockwiseSumReduce =
-        PartitionedBlockwiseReduction<AccDataType,
+    using BlockwiseSumReduce = PartitionedBlockwiseReduction<AccDataType,
                                                             BlockSize,
                                                             ThreadClusterLengths_M_K,
                                                             ThreadClusterArrangeOrder,
                                                             reduce::Add,
-                                      false, // ignored
-                                      detail::AccumulateWithNanIgnore<reduce::Add, AccDataType>>;
+                                                             true>;

-    using ThreadwiseSumReduce =
-        ThreadwiseReduction<AccDataType,
+    using ThreadwiseSumReduce = ThreadwiseReduction<AccDataType,
                                                    ThreadReduceSrcDesc_M_K,
                                                    ThreadReduceDstDesc_M,
                                                    reduce::Add,
-                            false, // ignored
-                            detail::AccumulateWithNanIgnore<reduce::Add, AccDataType>>;
+                                                    true>;

    static constexpr auto I0 = Number<0>{};
    static constexpr auto I1 = Number<1>{};