Skip to content

Commit

Permalink
[circlechef] Support mx data type
Browse files Browse the repository at this point in the history
This introduces MXFP4 and MXINT8.

ONE-DCO-1.0-Signed-off-by: Hyukjin Jeong <[email protected]>
  • Loading branch information
jinevening committed Nov 5, 2024
1 parent 2b89cb3 commit ba7eb39
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 0 deletions.
4 changes: 4 additions & 0 deletions compiler/circlechef/circle/src/Convert.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@ circlechef::TensorType as_circlechef_type(const circle::TensorType type)
return circlechef::INT16;
case circle::TensorType_INT4:
return circlechef::INT4;
case circle::TensorType_MXFP4:
return circlechef::MXFP4;
case circle::TensorType_MXINT8:
return circlechef::MXINT8;
// TODO handle other types
// TensorType_FLOAT16
// TensorType_STRING
Expand Down
4 changes: 4 additions & 0 deletions compiler/circlechef/core/src/Convert.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,10 @@ circle::TensorType as_circle_tensortype(const circlechef::TensorType &value)
return circle::TensorType_STRING;
case circlechef::BOOL:
return circle::TensorType_BOOL;
case circlechef::MXFP4:
return circle::TensorType_MXFP4;
case circlechef::MXINT8:
return circle::TensorType_MXINT8;
default:
break;
}
Expand Down
2 changes: 2 additions & 0 deletions compiler/circlechef/core/src/Convert.test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ TEST(ConvertTest, as_circle_tensortype)
ASSERT_EQ(circle::TensorType_UINT8, as_circle_tensortype(circlechef::UINT8));
ASSERT_EQ(circle::TensorType_UINT4, as_circle_tensortype(circlechef::UINT4));
ASSERT_EQ(circle::TensorType_BOOL, as_circle_tensortype(circlechef::BOOL));
ASSERT_EQ(circle::TensorType_MXFP4, as_circle_tensortype(circlechef::MXFP4));
ASSERT_EQ(circle::TensorType_MXINT8, as_circle_tensortype(circlechef::MXINT8));
}

TEST(ConvertTest, as_circle_tensortype_NEG)
Expand Down
2 changes: 2 additions & 0 deletions compiler/circlechef/proto/circlechef.proto
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ package circlechef;

// This enum value corresponds to TensorType in TensorFlow Lite schema
enum TensorType {
MXINT8 = -7;
MXFP4 = -6;
UINT4 = -1;
FLOAT32 = 0;
INT32 = 2;
Expand Down

0 comments on commit ba7eb39

Please sign in to comment.