From ad45cf8613211a11308dcae5637c829e59710668 Mon Sep 17 00:00:00 2001
From: Po Yen Chen <PoYen.Chen@amd.com>
Date: Tue, 9 Apr 2024 07:41:30 +0000
Subject: [PATCH] Support heterogeneous argument for binary function types

---
 include/ck_tile/core/numeric/math.hpp | 78 +++++++++++++++++++++------
 1 file changed, 63 insertions(+), 15 deletions(-)

diff --git a/include/ck_tile/core/numeric/math.hpp b/include/ck_tile/core/numeric/math.hpp
index 0c67a640af..868fff05b7 100644
--- a/include/ck_tile/core/numeric/math.hpp
+++ b/include/ck_tile/core/numeric/math.hpp
@@ -19,27 +19,75 @@ struct scales
     CK_TILE_HOST_DEVICE constexpr T operator()(T a) const { return s * a; }
 };
 
-template <typename T>
+template <typename Left = void, typename Right = Left>
 struct plus
 {
-    CK_TILE_HOST_DEVICE constexpr T operator()(T a, T b) const { return a + b; }
-};
-
-template <typename T>
-struct minus
-{
-    CK_TILE_HOST_DEVICE constexpr T operator()(T a, T b) const { return a - b; }
-};
-
-struct multiplies
-{
-    template <typename A, typename B>
-    CK_TILE_HOST_DEVICE constexpr auto operator()(const A& a, const B& b) const
+    CK_TILE_HOST_DEVICE constexpr auto operator()(const Left& lhs, const Right& rhs) const
+        -> decltype(lhs + rhs)
     {
-        return a * b;
+        return lhs + rhs;
     }
 };
 
+template <>
+struct plus<void, void>
+{
+    template <typename Left, typename Right>
+    CK_TILE_HOST_DEVICE constexpr auto operator()(const Left& lhs, const Right& rhs) const
+        -> decltype(lhs + rhs)
+    {
+        return lhs + rhs;
+    }
+};
+
+__host__ __device__ plus()->plus<void, void>;
+
+template <typename Left = void, typename Right = Left>
+struct minus
+{
+    CK_TILE_HOST_DEVICE constexpr auto operator()(const Left& lhs, const Right& rhs) const
+        -> decltype(lhs - rhs)
+    {
+        return lhs - rhs;
+    }
+};
+
+template <>
+struct minus<void, void>
+{
+    template <typename Left, typename Right>
+    CK_TILE_HOST_DEVICE constexpr auto operator()(const Left& lhs, const Right& rhs) const
+        -> decltype(lhs - rhs)
+    {
+        return lhs - rhs;
+    }
+};
+
+__host__ __device__ minus()->minus<void, void>;
+
+template <typename Left = void, typename Right = Left>
+struct multiplies
+{
+    CK_TILE_HOST_DEVICE constexpr auto operator()(const Left& lhs, const Right& rhs) const
+        -> decltype(lhs * rhs)
+    {
+        return lhs * rhs;
+    }
+};
+
+template <>
+struct multiplies<void, void>
+{
+    template <typename Left, typename Right>
+    CK_TILE_HOST_DEVICE constexpr auto operator()(const Left& lhs, const Right& rhs) const
+        -> decltype(lhs * rhs)
+    {
+        return lhs * rhs;
+    }
+};
+
+__host__ __device__ multiplies()->multiplies<void, void>;
+
 template <typename T>
 struct maximize
 {