Skip to content

Commit fa1d697

Browse files
committed
feat: add distinct bool data type to InfiniRT
1 parent 6b256e0 commit fa1d697

2 files changed

Lines changed: 23 additions & 4 deletions

File tree

‎src/data_type.h‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,11 @@ enum class DataType : std::int8_t {
2323
kFloat16,
2424
kBFloat16,
2525
kFloat32,
26-
kFloat64
26+
kFloat64,
27+
kBool
2728
};
2829

29-
constexpr ConstexprMap<DataType, std::size_t, 12> kDataTypeToSize{{{
30+
constexpr ConstexprMap<DataType, std::size_t, 13> kDataTypeToSize{{{
3031
{DataType::kInt8, 1},
3132
{DataType::kInt16, 2},
3233
{DataType::kInt32, 4},
@@ -39,9 +40,10 @@ constexpr ConstexprMap<DataType, std::size_t, 12> kDataTypeToSize{{{
3940
{DataType::kBFloat16, 2},
4041
{DataType::kFloat32, 4},
4142
{DataType::kFloat64, 8},
43+
{DataType::kBool, 1},
4244
}}};
4345

44-
constexpr ConstexprMap<DataType, std::string_view, 12> kDataTypeToDesc{{{
46+
constexpr ConstexprMap<DataType, std::string_view, 13> kDataTypeToDesc{{{
4547
{DataType::kInt8, "int8"},
4648
{DataType::kInt16, "int16"},
4749
{DataType::kInt32, "int32"},
@@ -54,9 +56,10 @@ constexpr ConstexprMap<DataType, std::string_view, 12> kDataTypeToDesc{{{
5456
{DataType::kBFloat16, "bfloat16"},
5557
{DataType::kFloat32, "float32"},
5658
{DataType::kFloat64, "float64"},
59+
{DataType::kBool, "bool"},
5760
}}};
5861

59-
constexpr ConstexprMap<std::string_view, DataType, 12> kStringToDataType{{{
62+
constexpr ConstexprMap<std::string_view, DataType, 13> kStringToDataType{{{
6063
{"int8", DataType::kInt8},
6164
{"int16", DataType::kInt16},
6265
{"int32", DataType::kInt32},
@@ -69,6 +72,7 @@ constexpr ConstexprMap<std::string_view, DataType, 12> kStringToDataType{{{
6972
{"bfloat16", DataType::kBFloat16},
7073
{"float32", DataType::kFloat32},
7174
{"float64", DataType::kFloat64},
75+
{"bool", DataType::kBool},
7276
}}};
7377

7478
struct Float16 {
@@ -164,6 +168,7 @@ using TypeMapType = typename TypeMap<dev, dtype>::type;
164168
};
165169

166170
DEFINE_DATA_TYPE_MAPPING(kUInt8, std::uint8_t)
171+
DEFINE_DATA_TYPE_MAPPING(kBool, bool)
167172
DEFINE_DATA_TYPE_MAPPING(kInt8, std::int8_t)
168173
DEFINE_DATA_TYPE_MAPPING(kUInt16, std::uint16_t)
169174
DEFINE_DATA_TYPE_MAPPING(kInt16, std::int16_t)

‎tests/test_core.cc‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,20 @@ void TestDataType(infini::rt::test::TestContext* context) {
159159
"float64", "float64 should have a stable name.");
160160
context->ExpectEqual(infini::rt::kStringToDataType.at("uint16"),
161161
DataType::kUInt16, "uint16 should parse by name.");
162+
context->Expect(DataType::kBool != DataType::kUInt8,
163+
"bool must be distinct from uint8.");
164+
context->ExpectEqual(infini::rt::kDataTypeToSize.at(DataType::kBool),
165+
std::size_t{1}, "bool should use one byte.");
166+
context->ExpectEqual(infini::rt::kDataTypeToDesc.at(DataType::kBool), "bool",
167+
"bool should have a stable name.");
168+
context->ExpectEqual(infini::rt::kStringToDataType.at("bool"),
169+
DataType::kBool, "bool should parse by name.");
170+
context->ExpectEqual(infini::rt::kStringToDataType.at("uint8"),
171+
DataType::kUInt8, "uint8 should remain distinct.");
172+
context->Expect(
173+
std::is_same_v<
174+
infini::rt::TypeMapType<Device::Type::kCpu, DataType::kBool>, bool>,
175+
"CPU bool should map to C++ bool.");
162176
}
163177

164178
void TestTensorViewRanks(infini::rt::test::TestContext* context) {

0 commit comments

Comments
 (0)