[RTG] Add set union operation (#7916)

llvm · Dec 9, 2024 · 2aecdf7 · 2aecdf7
1 parent de64a10
commit 2aecdf7
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 3 deletions.
diff --git a/include/circt/Dialect/RTG/IR/RTGOps.td b/include/circt/Dialect/RTG/IR/RTGOps.td
@@ -16,6 +16,7 @@ include "mlir/IR/Properties.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "circt/Dialect/RTG/IR/RTGInterfaces.td"
 
 // Base class for the operation in this dialect.
@@ -141,6 +142,23 @@ def SetDifferenceOp : RTGOp<"set_difference", [
   }];
 }
 
+def SetUnionOp : RTGOp<"set_union", [
+  Pure, SameOperandsAndResultType, Commutative
+]> {
+  let summary = "computes the union of sets";
+  let description = [{
+    Computes the union of the given sets. The list of sets must contain at
+    least one element.
+  }];
+
+  let arguments = (ins Variadic<SetType>:$sets);
+  let results = (outs SetType:$result);
+
+  let assemblyFormat = [{
+    $sets `:` qualified(type($result)) attr-dict
+  }];
+}
+
 //===- Bag Operations ------------------------------------------------------===//
 
 def BagCreateOp : RTGOp<"bag_create", [Pure, SameVariadicOperandSize]> {

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

diff --git a/lib/Dialect/RTG/Transforms/ElaborationPass.cpp b/lib/Dialect/RTG/Transforms/ElaborationPass.cpp
@@ -514,6 +514,17 @@ class Elaborator : public RTGOpVisitor<Elaborator, FailureOr<DeletionKind>,
     return DeletionKind::Delete;
   }
 
+  FailureOr<DeletionKind>
+  visitOp(SetUnionOp op, function_ref<void(Operation *)> addToWorklist) {
+    SetVector<ElaboratorValue *> result;
+    for (auto set : op.getSets())
+      result.set_union(cast<SetValue>(state.at(set))->getSet());
+
+    internalizeResult<SetValue>(op.getResult(), std::move(result),
+                                op.getType());
+    return DeletionKind::Delete;
+  }
+
   FailureOr<DeletionKind>
   visitOp(BagCreateOp op, function_ref<void(Operation *)> addToWorklist) {
     MapVector<ElaboratorValue *, uint64_t> bag;

diff --git a/test/Dialect/RTG/IR/basic.mlir b/test/Dialect/RTG/IR/basic.mlir
@@ -30,11 +30,13 @@ func.func @sets(%arg0: i32, %arg1: i32) {
   // CHECK: [[SET:%.+]] = rtg.set_create %arg0, %arg1 : i32
   // CHECK: [[R:%.+]] = rtg.set_select_random [[SET]] : !rtg.set<i32>
   // CHECK: [[EMPTY:%.+]] = rtg.set_create : i32
-  // CHECK: rtg.set_difference [[SET]], [[EMPTY]] : !rtg.set<i32>
+  // CHECK: [[DIFF:%.+]] = rtg.set_difference [[SET]], [[EMPTY]] : !rtg.set<i32>
+  // CHECK: rtg.set_union [[SET]], [[DIFF]] : !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>
 
   return
 }

diff --git a/test/Dialect/RTG/IR/errors.mlir b/test/Dialect/RTG/IR/errors.mlir
@@ -66,3 +66,10 @@ rtg.sequence @seq {
   // expected-error @below {{operand types must match bag element type}}
   "rtg.bag_create"(%arg0, %arg1, %arg2, %arg2){} : (i64, i64, index, index) -> !rtg.bag<i32>
 }
+
+// -----
+
+rtg.sequence @seq {
+  // expected-error @below {{expected 1 or more operands, but found 0}}
+  rtg.set_union : !rtg.set<i32>
+}
diff --git a/test/Dialect/RTG/Transform/elaboration.mlir b/test/Dialect/RTG/Transform/elaboration.mlir
@@ -18,7 +18,9 @@ rtg.test @setOperations : !rtg.dict<> {
   %1 = arith.constant 3 : i32
   %2 = arith.constant 4 : i32
   %3 = arith.constant 5 : i32
-  %set = rtg.set_create %0, %1, %2, %0 : i32
+  %set0 = rtg.set_create %0, %1, %0 : i32
+  %set1 = rtg.set_create %2, %0 : i32
+  %set = rtg.set_union %set0, %set1 : !rtg.set<i32>
   %4 = rtg.set_select_random %set : !rtg.set<i32> {rtg.elaboration_custom_seed = 1}
   %new_set = rtg.set_create %3, %4 : i32
   %diff = rtg.set_difference %set, %new_set : !rtg.set<i32>