From 323c8170bffdd11d774437b450e42d842e203517 Mon Sep 17 00:00:00 2001
From: Tri Dao <tridao@users.noreply.github.com>
Date: Thu, 5 Sep 2024 20:25:03 -0700
Subject: [PATCH] Support ComputeFn where output type differs from input type
 (#1771)

This is useful for e.g. function taking in 2 float inputs and turn them to complex
---
 .../sm90_visitor_compute_tma_warpspecialized.hpp       | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp b/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
index 0b12badc7d..8f5ceb5489 100644
--- a/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
+++ b/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
@@ -181,14 +181,20 @@ struct Sm90Compute {
         },
         [&] (auto&&... cvt_frg_inputs) {
           using ComputeOutput = ComputeFn<Array<ElementCompute, FragmentSize>>;
-          using ConvertOutput = NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize, RoundStyle>;
           ComputeOutput compute_output{};
-          ConvertOutput convert_output{};
 
           if constexpr (cute::is_same_v<Arguments, EmptyArguments>) {
+            using ElementComputeOutput =
+                typename cute::remove_cvref_t<decltype(compute_output(cvt_frg_inputs...))>::Element;
+            using ConvertOutput = NumericArrayConverter<ElementOutput, ElementComputeOutput, FragmentSize, RoundStyle>;
+            ConvertOutput convert_output{};
             return convert_output(compute_output(cvt_frg_inputs...));
           }
           else {
+            using ElementComputeOutput =
+                typename cute::remove_cvref_t<decltype(compute_output(cvt_frg_inputs..., params))>::Element;
+            using ConvertOutput = NumericArrayConverter<ElementOutput, ElementComputeOutput, FragmentSize, RoundStyle>;
+            ConvertOutput convert_output{};
             return convert_output(compute_output(cvt_frg_inputs..., params));
           }
         }