[RTG] Add set_convert_to_bag operation

frontends/PyRTG/src/pyrtg/sets.py

@@ -91,6 +91,13 @@ def cartesian_product(*args: Set) -> Set:
     assert len(args) > 0, "at least one set must be provided"
     return rtg.SetCartesianProductOp(args)
 
+  def to_bag(self) -> Value:
+    """
+    Returns this set converted to a bag. Does not modify this object.
+    """
+
+    return rtg.SetConvertToBagOp(self)
+
   def get_random(self) -> Value:
     """
     Returns an element from the set picked uniformly at random. If the set is

frontends/PyRTG/test/basic.py

@@ -438,10 +438,12 @@ 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>
+# MLIR: rtg.set_convert_to_bag %a : !rtg.set<index>
 
 
 @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())
+  int_consumer(a.to_bag().get_random())

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

@@ -390,6 +390,24 @@ def SetCartesianProductOp : RTGOp<"set_cartesian_product", [
   let assemblyFormat = "$inputs `:` qualified(type($inputs)) attr-dict";
 }
 
+def SetConvertToBagOp : RTGOp<"set_convert_to_bag", [
+  Pure,
+  TypesMatchWith<"element type of set must match the bag's element type",
+                 "input", "result",
+                 "BagType::get(cast<SetType>($_self).getElementType())">,
+]> {
+  let summary = "convert a set to a bag";
+  let description = [{
+    This operation converts a set to a bag. Each element in the set occurs
+    exactly once in the resulting bag.
+  }];
+
+  let arguments = (ins SetType:$input);
+  let results = (outs BagType:$result);
+
+  let assemblyFormat = "$input `:` qualified(type($input)) attr-dict";
+}
+
 //===- Bag Operations ------------------------------------------------------===//
 
 def BagCreateOp : RTGOp<"bag_create", [Pure, SameVariadicOperandSize]> {

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

@@ -90,6 +90,11 @@ def BagType : RTGTypeDef<"Bag"> {
 
   let mnemonic = "bag";
   let assemblyFormat = "`<` $elementType `>`";
+
+  let builders = [
+    TypeBuilderWithInferredContext<(ins "::mlir::Type":$elementType),
+      "return $_get(elementType.getContext(), elementType);">,
+  ];
 }
 
 class SetTypeOf<Type elementType> : ContainerType<

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

@@ -52,7 +52,7 @@ class RTGOpVisitor {
             RandomizeSequenceOp, EmbedSequenceOp, InterleaveSequencesOp,
             // Sets
             SetCreateOp, SetSelectRandomOp, SetDifferenceOp, SetUnionOp,
-            SetSizeOp, SetCartesianProductOp,
+            SetSizeOp, SetCartesianProductOp, SetConvertToBagOp,
             // Arrays
             ArrayCreateOp, ArrayExtractOp, ArrayInjectOp, ArraySizeOp,
             // Tuples
@@ -110,6 +110,7 @@ class RTGOpVisitor {
   HANDLE(SetUnionOp, Unhandled);
   HANDLE(SetSizeOp, Unhandled);
   HANDLE(SetCartesianProductOp, Unhandled);
+  HANDLE(SetConvertToBagOp, Unhandled);
   HANDLE(BagCreateOp, Unhandled);
   HANDLE(BagSelectRandomOp, Unhandled);
   HANDLE(BagDifferenceOp, Unhandled);

lib/Dialect/RTG/Transforms/ElaborationPass.cpp

@@ -1173,6 +1173,16 @@ class Elaborator : public RTGOpVisitor<Elaborator, FailureOr<DeletionKind>> {
     return DeletionKind::Delete;
   }
 
+  FailureOr<DeletionKind> visitOp(SetConvertToBagOp op) {
+    auto set = get<SetStorage *>(op.getInput())->set;
+    MapVector<ElaboratorValue, uint64_t> bag;
+    for (auto val : set)
+      bag.insert({val, 1});
+    state[op.getResult()] = sharedState.internalizer.internalize<BagStorage>(
+        std::move(bag), op.getType());
+    return DeletionKind::Delete;
+  }
+
   FailureOr<DeletionKind> visitOp(BagCreateOp op) {
     MapVector<ElaboratorValue, uint64_t> bag;
     for (auto [val, multiple] :

test/Dialect/RTG/IR/basic.mlir

@@ -67,13 +67,15 @@ func.func @sets(%arg0: i32, %arg1: i32) -> !rtg.set<tuple<i32, i32>> {
   // CHECK: rtg.set_union [[SET]], [[DIFF]] : !rtg.set<i32>
   // CHECK: rtg.set_size [[SET]] : !rtg.set<i32>
   // CHECK: rtg.set_cartesian_product [[SET]], [[SET]] : !rtg.set<i32>, !rtg.set<i32>
+  // CHECK: rtg.set_convert_to_bag [[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>
   %prod = rtg.set_cartesian_product %set, %set : !rtg.set<i32>, !rtg.set<i32>
+  %bag = rtg.set_convert_to_bag %set : !rtg.set<i32>
 
   return %prod : !rtg.set<tuple<i32, i32>>
 }

test/Dialect/RTG/Transform/elaboration.mlir

@@ -11,6 +11,7 @@ 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}
+func.func @dummy12(%arg0: !rtg.bag<index>) -> () {return}
 
 // CHECK-LABEL: @immediates
 rtg.test @immediates() {
@@ -34,7 +35,6 @@ rtg.test @setOperations() {
   // CHECK-NEXT: [[V2:%.+]] = index.constant 4
   // CHECK-NEXT: [[V3:%.+]] = rtg.set_create [[V1]], [[V2]] : index
   // CHECK-NEXT: func.call @dummy1([[V0]], [[V1]], [[V3]]) :
-  // CHECK-NEXT: }
   %0 = index.constant 2
   %1 = index.constant 3
   %2 = index.constant 4
@@ -47,6 +47,12 @@ rtg.test @setOperations() {
   %diff = rtg.set_difference %set, %new_set : !rtg.set<index>
   %5 = rtg.set_select_random %diff : !rtg.set<index> {rtg.elaboration_custom_seed = 2}
   func.call @dummy1(%4, %5, %diff) : (index, index, !rtg.set<index>) -> ()
+
+  // CHECK-NEXT: [[V4:%.+]] = index.constant 1 
+  // CHECK-NEXT: [[V5:%.+]] = rtg.bag_create ([[V4]] x [[V2]], [[V4]] x [[V0]]) : index 
+  // CHECK-NEXT: func.call @dummy12([[V5]]) : (!rtg.bag<index>)
+  %6 = rtg.set_convert_to_bag %set1 : !rtg.set<index>
+  func.call @dummy12(%6) : (!rtg.bag<index>) -> ()
 }
 
 // CHECK-LABEL: rtg.test @setCartesianProduct