[RTG] Add bag_convert_to_set operation

frontends/PyRTG/src/pyrtg/bags.py

@@ -129,6 +129,14 @@ def get_random_and_exclude(self) -> Value:
     self._value = self.exclude(r)._get_ssa_value()
     return r
 
+  def to_set(self) -> Value:
+    """
+    Returns this bag converted to a set, i.e., all duplicates are dropped. Does
+    not modify this object.
+    """
+
+    return rtg.BagConvertToSetOp(self)
+
   def _get_ssa_value(self) -> ir.Value:
     return self._value
 

frontends/PyRTG/test/basic.py

@@ -437,8 +437,11 @@ def test90_tuples(a, b, tup):
 
 # MLIR-LABEL: rtg.test @test91_sets
 # MLIR-NEXT: rtg.set_cartesian_product %a, %b : !rtg.set<index>, !rtg.set<i1>
+# MLIR: rtg.bag_convert_to_set %c : !rtg.bag<index>
 
 
-@test(("a", Set.type(Integer.type())), ("b", Set.type(Bool.type())))
-def test91_sets(a, b):
+@test(("a", Set.type(Integer.type())), ("b", Set.type(Bool.type())),
+      ("c", Bag.type(Integer.type())))
+def test91_sets(a, b, c):
   seq2(Set.cartesian_product(a, b))
+  int_consumer(c.to_set().get_random())

include/circt/Dialect/RTG/IR/RTGOps.td

@@ -487,6 +487,24 @@ def BagUniqueSizeOp : RTGOp<"bag_unique_size", [Pure]> {
   }];
 }
 
+def BagConvertToSetOp : RTGOp<"bag_convert_to_set", [
+  Pure,
+  TypesMatchWith<"element type of set must match the bag's element type",
+                 "input", "result",
+                 "SetType::get(cast<BagType>($_self).getElementType())">,
+]> {
+  let summary = "convert a bag to a set";
+  let description = [{
+    This operation converts a bag to a set by dropping all duplicate elements.
+    For example, the bag `{a, a, b}` is converted to `{a, b}`.
+  }];
+
+  let arguments = (ins BagType:$input);
+  let results = (outs SetType:$result);
+
+  let assemblyFormat = "$input `:` qualified(type($input)) attr-dict";
+}
+
 //===- Array Operations -------------------------------------------------===//
 
 def ArrayCreateOp : RTGOp<"array_create", [

include/circt/Dialect/RTG/IR/RTGVisitors.h

@@ -36,7 +36,7 @@ class RTGOpVisitor {
             ConstantOp,
             // Bags
             BagCreateOp, BagSelectRandomOp, BagDifferenceOp, BagUnionOp,
-            BagUniqueSizeOp,
+            BagUniqueSizeOp, BagConvertToSetOp,
             // Contexts
             OnContextOp, ContextSwitchOp,
             // Labels
@@ -115,6 +115,7 @@ class RTGOpVisitor {
   HANDLE(BagDifferenceOp, Unhandled);
   HANDLE(BagUnionOp, Unhandled);
   HANDLE(BagUniqueSizeOp, Unhandled);
+  HANDLE(BagConvertToSetOp, Unhandled);
   HANDLE(ArrayCreateOp, Unhandled);
   HANDLE(ArrayExtractOp, Unhandled);
   HANDLE(ArrayInjectOp, Unhandled);

lib/Dialect/RTG/Transforms/ElaborationPass.cpp

@@ -1264,6 +1264,16 @@ class Elaborator : public RTGOpVisitor<Elaborator, FailureOr<DeletionKind>> {
     return DeletionKind::Delete;
   }
 
+  FailureOr<DeletionKind> visitOp(BagConvertToSetOp op) {
+    auto bag = get<BagStorage *>(op.getInput())->bag;
+    SetVector<ElaboratorValue> set;
+    for (auto [k, v] : bag)
+      set.insert(k);
+    state[op.getResult()] = sharedState.internalizer.internalize<SetStorage>(
+        std::move(set), op.getType());
+    return DeletionKind::Delete;
+  }
+
   FailureOr<DeletionKind> visitOp(FixedRegisterOp op) {
     return visitPureOp(op);
   }

test/Dialect/RTG/IR/basic.mlir

@@ -87,13 +87,15 @@ rtg.sequence @bags(%arg0: i32, %arg1: i32, %arg2: index) {
   // 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>
+  // CHECK: rtg.bag_convert_to_set [[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>
+  %set = rtg.bag_convert_to_set %bag : !rtg.bag<i32>
 }
 
 // CHECK-LABEL: rtg.target @empty_target : !rtg.dict<> {

test/Dialect/RTG/Transform/elaboration.mlir

@@ -10,6 +10,7 @@ func.func @dummy7(%arg0: !rtg.array<index>) -> () {return}
 func.func @dummy8(%arg0: tuple<index, index>) -> () {return}
 func.func @dummy9(%arg0: !rtg.set<tuple<index, i1, !rtgtest.ireg>>) -> () {return}
 func.func @dummy10(%arg0: !rtg.set<tuple<index>>) -> () {return}
+func.func @dummy11(%arg0: !rtg.set<index>) -> () {return}
 
 // CHECK-LABEL: @immediates
 rtg.test @immediates() {
@@ -117,6 +118,11 @@ rtg.test @bagOperations() {
   %diff2 = rtg.bag_difference %bag, %new_bag inf : !rtg.bag<index>
   %4 = rtg.bag_select_random %diff2 : !rtg.bag<index> {rtg.elaboration_custom_seed = 5}
   func.call @dummy4(%3, %4, %diff, %diff2) : (index, index, !rtg.bag<index>, !rtg.bag<index>) -> ()
+
+  // CHECK-NEXT: [[SET:%.+]] = rtg.set_create [[V0]], [[V2]] :
+  // CHECK-NEXT: func.call @dummy11([[SET]])
+  %5 = rtg.bag_convert_to_set %bag0 : !rtg.bag<index>
+  func.call @dummy11(%5) : (!rtg.set<index>) -> ()
 }
 
 // CHECK-LABEL: rtg.test @setSize