[RTG] Add set_size op for sets and bag_unique_size op for bags (#7920)

llvm · Dec 9, 2024 · 7f8c80b · 7f8c80b
1 parent daf1bda
commit 7f8c80b
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 3 deletions.
diff --git a/include/circt/Dialect/RTG/IR/RTGOps.td b/include/circt/Dialect/RTG/IR/RTGOps.td
@@ -159,6 +159,17 @@ def SetUnionOp : RTGOp<"set_union", [
   }];
 }
 
+def SetSizeOp : RTGOp<"set_size", [Pure]> {
+  let summary = "returns the number of elements in the set";
+
+  let arguments = (ins SetType:$set);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = [{
+    $set `:` qualified(type($set)) attr-dict
+  }];
+}
+
 //===- Bag Operations ------------------------------------------------------===//
 
 def BagCreateOp : RTGOp<"bag_create", [Pure, SameVariadicOperandSize]> {
@@ -241,6 +252,21 @@ def BagUnionOp : RTGOp<"bag_union", [
   }];
 }
 
+def BagUniqueSizeOp : RTGOp<"bag_unique_size", [Pure]> {
+  let summary = "returns the number of unique elements in the bag";
+  let description = [{
+    This operation returns the number of unique elements in the bag, i.e., for
+    the bag `{a, a, b, c, c}` it returns 3.
+  }];
+
+  let arguments = (ins BagType:$bag);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = [{
+    $bag `:` qualified(type($bag)) attr-dict
+  }];
+}
+
 //===- Test Specification Operations --------------------------------------===//
 
 def TestOp : RTGOp<"test", [

diff --git a/include/circt/Dialect/RTG/IR/RTGVisitors.h b/include/circt/Dialect/RTG/IR/RTGVisitors.h
@@ -32,9 +32,10 @@ class RTGOpVisitor {
     auto *thisCast = static_cast<ConcreteType *>(this);
     return TypeSwitch<Operation *, ResultType>(op)
         .template Case<SequenceOp, SequenceClosureOp, SetCreateOp,
-                       SetSelectRandomOp, SetDifferenceOp, SetUnionOp, TestOp,
-                       InvokeSequenceOp, BagCreateOp, BagSelectRandomOp,
-                       BagDifferenceOp, BagUnionOp, TargetOp, YieldOp>(
+                       SetSelectRandomOp, SetDifferenceOp, SetUnionOp,
+                       SetSizeOp, TestOp, InvokeSequenceOp, BagCreateOp,
+                       BagSelectRandomOp, BagDifferenceOp, BagUnionOp,
+                       BagUniqueSizeOp, TargetOp, YieldOp>(
             [&](auto expr) -> ResultType {
               return thisCast->visitOp(expr, args...);
             })
@@ -90,10 +91,12 @@ class RTGOpVisitor {
   HANDLE(SetSelectRandomOp, Unhandled);
   HANDLE(SetDifferenceOp, Unhandled);
   HANDLE(SetUnionOp, Unhandled);
+  HANDLE(SetSizeOp, Unhandled);
   HANDLE(BagCreateOp, Unhandled);
   HANDLE(BagSelectRandomOp, Unhandled);
   HANDLE(BagDifferenceOp, Unhandled);
   HANDLE(BagUnionOp, Unhandled);
+  HANDLE(BagUniqueSizeOp, Unhandled);
   HANDLE(TestOp, Unhandled);
   HANDLE(TargetOp, Unhandled);
   HANDLE(YieldOp, Unhandled);

diff --git a/test/Dialect/RTG/IR/basic.mlir b/test/Dialect/RTG/IR/basic.mlir
@@ -32,11 +32,13 @@ func.func @sets(%arg0: i32, %arg1: i32) {
   // CHECK: [[EMPTY:%.+]] = rtg.set_create : i32
   // CHECK: [[DIFF:%.+]] = rtg.set_difference [[SET]], [[EMPTY]] : !rtg.set<i32>
   // CHECK: rtg.set_union [[SET]], [[DIFF]] : !rtg.set<i32>
+  // CHECK: rtg.set_size [[SET]] : !rtg.set<i32>
   %set = rtg.set_create %arg0, %arg1 : i32
   %r = rtg.set_select_random %set : !rtg.set<i32>
   %empty = rtg.set_create : i32
   %diff = rtg.set_difference %set, %empty : !rtg.set<i32>
   %union = rtg.set_union %set, %diff : !rtg.set<i32>
+  %size = rtg.set_size %set : !rtg.set<i32>
 
   return
 }
@@ -50,12 +52,14 @@ rtg.sequence @bags {
   // CHECK: [[DIFF:%.+]] = rtg.bag_difference [[BAG]], [[EMPTY]] : !rtg.bag<i32> {rtg.some_attr}
   // CHECK: rtg.bag_difference [[BAG]], [[EMPTY]] inf : !rtg.bag<i32>
   // CHECK: rtg.bag_union [[BAG]], [[EMPTY]], [[DIFF]] : !rtg.bag<i32>
+  // CHECK: rtg.bag_unique_size [[BAG]] : !rtg.bag<i32>
   %bag = rtg.bag_create (%arg2 x %arg0, %arg2 x %arg1) : i32 {rtg.some_attr}
   %r = rtg.bag_select_random %bag : !rtg.bag<i32> {rtg.some_attr}
   %empty = rtg.bag_create : i32
   %diff = rtg.bag_difference %bag, %empty : !rtg.bag<i32> {rtg.some_attr}
   %diff2 = rtg.bag_difference %bag, %empty inf : !rtg.bag<i32>
   %union = rtg.bag_union %bag, %empty, %diff : !rtg.bag<i32>
+  %size = rtg.bag_unique_size %bag : !rtg.bag<i32>
 }
 
 // CHECK-LABEL: rtg.target @empty_target : !rtg.dict<> {