support data type w2 used in cutlass_library (#1517)

This commit is contained in:
chenwei 2024-08-06 23:15:18 +08:00 committed by GitHub
parent 19b4c5e065
commit e22ba590cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -69,6 +69,7 @@ GeneratorTargetNames = {
class DataType(enum.Enum): class DataType(enum.Enum):
void = enum_auto() # primarily used to disable C tensor for epilogues void = enum_auto() # primarily used to disable C tensor for epilogues
b1 = enum_auto() b1 = enum_auto()
u2 = enum_auto()
u4 = enum_auto() u4 = enum_auto()
u8 = enum_auto() u8 = enum_auto()
u16 = enum_auto() u16 = enum_auto()
@ -119,6 +120,7 @@ ShortDataTypeNames = {
DataTypeNames = { DataTypeNames = {
DataType.void: "void", DataType.void: "void",
DataType.b1: "b1", DataType.b1: "b1",
DataType.u2: "u2",
DataType.u4: "u4", DataType.u4: "u4",
DataType.u8: "u8", DataType.u8: "u8",
DataType.u16: "u16", DataType.u16: "u16",
@ -156,6 +158,7 @@ DataTypeNames = {
DataTypeTag = { DataTypeTag = {
DataType.void: "void", DataType.void: "void",
DataType.b1: "cutlass::uint1b_t", DataType.b1: "cutlass::uint1b_t",
DataType.u2: "cutlass::uint2b_t",
DataType.u4: "cutlass::uint4b_t", DataType.u4: "cutlass::uint4b_t",
DataType.u8: "uint8_t", DataType.u8: "uint8_t",
DataType.u16: "uint16_t", DataType.u16: "uint16_t",
@ -193,6 +196,7 @@ DataTypeTag = {
DataTypeSize = { DataTypeSize = {
DataType.void: 0, DataType.void: 0,
DataType.b1: 1, DataType.b1: 1,
DataType.u2: 2,
DataType.u4: 4, DataType.u4: 4,
DataType.u8: 8, DataType.u8: 8,
DataType.u16: 16, DataType.u16: 16,