From 7f8c80b7784628860a891b57d5f689ad112b9513 Mon Sep 17 00:00:00 2001
From: Martin Erhart <martin.erhart@sifive.com>
Date: Mon, 9 Dec 2024 15:01:25 +0000
Subject: [PATCH] [RTG] Add set_size op for sets and bag_unique_size op for
 bags (#7920)

---
 include/circt/Dialect/RTG/IR/RTGOps.td     | 26 ++++++++++++++++++++++
 include/circt/Dialect/RTG/IR/RTGVisitors.h |  9 +++++---
 test/Dialect/RTG/IR/basic.mlir             |  4 ++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/include/circt/Dialect/RTG/IR/RTGOps.td b/include/circt/Dialect/RTG/IR/RTGOps.td
index 9f48655ec47a..943d9d0dbb36 100644
--- 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
index 0423bb72f699..4477a7d438d8 100644
--- 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
index ae49bbf9bb18..d3c29199a962 100644
--- 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<> {