[FIRRTL][CAPI] Add function for getting mask type

llvm · Apr 22, 2024 · 172b5ea · 172b5ea
1 parent db973f1
commit 172b5ea
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 0 deletions.
diff --git a/include/circt-c/Dialect/FIRRTL.h b/include/circt-c/Dialect/FIRRTL.h
@@ -114,6 +114,8 @@ MLIR_CAPI_EXPORTED MlirType
 firrtlTypeGetClass(MlirContext ctx, MlirAttribute name, size_t numberOfElements,
                    const FIRRTLClassElement *elements);
 
+MLIR_CAPI_EXPORTED MlirType firrtlTypeGetMaskType(MlirType type);
+
 //===----------------------------------------------------------------------===//
 // Attribute API.
 //===----------------------------------------------------------------------===//

diff --git a/lib/CAPI/Dialect/FIRRTL.cpp b/lib/CAPI/Dialect/FIRRTL.cpp
@@ -163,6 +163,12 @@ MlirType firrtlTypeGetClass(MlirContext ctx, MlirAttribute name,
   return wrap(ClassType::get(unwrap(ctx), nameSymbol, classElements));
 }
 
+MlirType firrtlTypeGetMaskType(MlirType type) {
+  auto baseType = type_dyn_cast<FIRRTLBaseType>(unwrap(type));
+  assert(baseType && "unexpected type, must be base type");
+  return wrap(baseType.getMaskType());
+}
+
 //===----------------------------------------------------------------------===//
 // Attribute API.
 //===----------------------------------------------------------------------===//

diff --git a/test/CAPI/firrtl.c b/test/CAPI/firrtl.c
@@ -21,6 +21,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+#define ARRAY_SIZE(arr) (sizeof((arr)) / sizeof((arr)[0]))
+
 void dumpCallback(MlirStringRef message, void *userData) {
   fprintf(stderr, "%.*s", (int)message.length, message.data);
 }
@@ -184,12 +186,60 @@ void testAttrGetIntegerFromString(MlirContext ctx) {
                       mlirStringRefCreateFromCString("114514"), 10));
 }
 
+void testTypeGetMaskType(MlirContext ctx) {
+  assert(mlirTypeEqual(firrtlTypeGetMaskType(firrtlTypeGetUInt(ctx, 32)),
+                       firrtlTypeGetUInt(ctx, 1)));
+  assert(mlirTypeEqual(firrtlTypeGetMaskType(firrtlTypeGetSInt(ctx, 64)),
+                       firrtlTypeGetUInt(ctx, 1)));
+
+  FIRRTLBundleField lhsFields[] = {
+      {
+          .name = mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("f1")),
+          .isFlip = false,
+          .type = firrtlTypeGetUInt(ctx, 32),
+      },
+      {
+          .name = mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("f2")),
+          .isFlip = false,
+          .type = firrtlTypeGetSInt(ctx, 64),
+      },
+      {
+          .name = mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("f3")),
+          .isFlip = false,
+          .type = firrtlTypeGetClock(ctx),
+      },
+  };
+  FIRRTLBundleField rhsFields[] = {
+      {
+          .name = mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("f1")),
+          .isFlip = false,
+          .type = firrtlTypeGetUInt(ctx, 1),
+      },
+      {
+          .name = mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("f2")),
+          .isFlip = false,
+          .type = firrtlTypeGetUInt(ctx, 1),
+      },
+      {
+          .name = mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("f3")),
+          .isFlip = false,
+          .type = firrtlTypeGetUInt(ctx, 1),
+      },
+  };
+  MlirType lhsBundle =
+      firrtlTypeGetBundle(ctx, ARRAY_SIZE(lhsFields), lhsFields);
+  MlirType rhsBundle =
+      firrtlTypeGetBundle(ctx, ARRAY_SIZE(rhsFields), rhsFields);
+  assert(mlirTypeEqual(firrtlTypeGetMaskType(lhsBundle), rhsBundle));
+}
+
 int main(void) {
   MlirContext ctx = mlirContextCreate();
   mlirDialectHandleLoadDialect(mlirGetDialectHandle__firrtl__(), ctx);
   testExport(ctx);
   testValueFoldFlow(ctx);
   testImportAnnotations(ctx);
   testAttrGetIntegerFromString(ctx);
+  testTypeGetMaskType(ctx);
   return 0;
 }