From e22ba590cd8a7eebea8f53c81b5740d905021654 Mon Sep 17 00:00:00 2001
From: chenwei <15601910741@163.com>
Date: Tue, 6 Aug 2024 23:15:18 +0800
Subject: [PATCH] support data type w2 used in cutlass_library (#1517)

---
 python/cutlass_library/library.py | 4 ++++
 1 file changed, 4 insertions(+)

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