From 0246565d15cf7790dffb2cae68b7f70c552a85f4 Mon Sep 17 00:00:00 2001
From: Sergey Kozub <skozub@nvidia.com>
Date: Tue, 17 Sep 2024 09:35:26 +0200
Subject: [PATCH] fix signbit

---
 ml_dtypes/include/mxfloat.h     | 13 +++++++++----
 ml_dtypes/tests/mxfloat_test.cc |  8 ++++++++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/ml_dtypes/include/mxfloat.h b/ml_dtypes/include/mxfloat.h
index 2741cecc..9a02862b 100644
--- a/ml_dtypes/include/mxfloat.h
+++ b/ml_dtypes/include/mxfloat.h
@@ -338,7 +338,7 @@ struct Traits<float4_e2m1fn>
 namespace Eigen {
 namespace numext {
 
-#define MXFLOAT_EIGEN_BITCAST_IMPL(Type)                                 \
+#define MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(Type)                     \
   template <>                                                            \
   EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint8_t bit_cast<uint8_t, Type>( \
       const Type& x) {                                                   \
@@ -348,11 +348,16 @@ namespace numext {
   EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Type bit_cast<Type, uint8_t>(    \
       const uint8_t& x) {                                                \
     return Type::FromRep(x);                                             \
+  }                                                                      \
+  template <>                                                            \
+  EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Type signbit(const Type& x) {    \
+    int8_t t = bit_cast<int8_t, Type>(x) << (8 - Type::kBits);           \
+    return bit_cast<Type, int8_t>(t >> 7);                               \
   }
 
-MXFLOAT_EIGEN_BITCAST_IMPL(ml_dtypes::float6_e2m3fn)
-MXFLOAT_EIGEN_BITCAST_IMPL(ml_dtypes::float6_e3m2fn)
-MXFLOAT_EIGEN_BITCAST_IMPL(ml_dtypes::float4_e2m1fn)
+MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(ml_dtypes::float6_e2m3fn)
+MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(ml_dtypes::float6_e3m2fn)
+MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(ml_dtypes::float4_e2m1fn)
 
 #undef MXFLOAT_EIGEN_BITCAST_IMPL
 
diff --git a/ml_dtypes/tests/mxfloat_test.cc b/ml_dtypes/tests/mxfloat_test.cc
index 834c1055..53fad2e0 100644
--- a/ml_dtypes/tests/mxfloat_test.cc
+++ b/ml_dtypes/tests/mxfloat_test.cc
@@ -125,6 +125,14 @@ TYPED_TEST(FloatMXTest, Negate) {
   }
 }
 
+TYPED_TEST(FloatMXTest, Signbit) {
+  using FloatMX = TypeParam;
+
+  FloatMX one(1.0);
+  EXPECT_EQ(Eigen::numext::signbit(one).rep(), 0x00);
+  EXPECT_EQ(Eigen::numext::signbit(-one).rep(), 0xff);
+}
+
 TYPED_TEST(FloatMXTest, BitCasts) {
   using FloatMX = TypeParam;