diff --git a/include/circt/Dialect/RTG/IR/RTGOps.td b/include/circt/Dialect/RTG/IR/RTGOps.td index 837418337a26..0387d6342983 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 SetGetSizeOp : RTGOp<"set_get_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 BagGetSizeOp : RTGOp<"bag_get_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..f21b0a12c423 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(this); return TypeSwitch(op) .template Case( + SetSelectRandomOp, SetDifferenceOp, SetUnionOp, + SetGetSizeOp, TestOp, InvokeSequenceOp, BagCreateOp, + BagSelectRandomOp, BagDifferenceOp, BagUnionOp, + BagGetSizeOp, 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(SetGetSizeOp, Unhandled); HANDLE(BagCreateOp, Unhandled); HANDLE(BagSelectRandomOp, Unhandled); HANDLE(BagDifferenceOp, Unhandled); HANDLE(BagUnionOp, Unhandled); + HANDLE(BagGetSizeOp, 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 c50b698ec3eb..bc344d35f53b 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 // CHECK: rtg.set_union [[SET]], [[DIFF]] : !rtg.set + // CHECK: rtg.set_get_size [[SET]] : !rtg.set %set = rtg.set_create %arg0, %arg1 : i32 %r = rtg.set_select_random %set : !rtg.set %empty = rtg.set_create : i32 %diff = rtg.set_difference %set, %empty : !rtg.set %union = rtg.set_union %set, %diff : !rtg.set + %size = rtg.set_get_size %set : !rtg.set return } @@ -50,12 +52,14 @@ rtg.sequence @bags { // CHECK: [[DIFF:%.+]] = rtg.bag_difference [[BAG]], [[EMPTY]] : !rtg.bag // CHECK: rtg.bag_difference [[BAG]], [[EMPTY]] inf : !rtg.bag // CHECK: rtg.bag_union [[BAG]], [[EMPTY]], [[DIFF]] : !rtg.bag + // CHECK: rtg.bag_get_size [[BAG]] : !rtg.bag %bag = rtg.bag_create (%arg2 x %arg0, %arg2 x %arg1) : i32 %r = rtg.bag_select_random %bag : !rtg.bag %empty = rtg.bag_create : i32 %diff = rtg.bag_difference %bag, %empty : !rtg.bag %diff2 = rtg.bag_difference %bag, %empty inf : !rtg.bag %union = rtg.bag_union %bag, %empty, %diff : !rtg.bag + %size = rtg.bag_get_size %bag : !rtg.bag } // CHECK-LABEL: rtg.target @empty_target : !rtg.dict<> {