[RTG] Add set union operation

llvm · Nov 28, 2024 · 2bf1d4d · 2bf1d4d
1 parent 0ae6ae2
commit 2bf1d4d
Show file tree

Hide file tree

Showing 6 changed files with 57 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. If 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
@@ -583,6 +583,30 @@ class Elaborator
     return DeletionKind::Delete;
   }
 
+  FailureOr<DeletionKind>
+  visitOp(SetUnionOp op, function_ref<void(Operation *)> addToWorklist) {
+    SmallVector<ElaboratorValue *> result;
+    SmallVector<ElaboratorValue *> debugMap;
+    for (auto set : op.getSets()) {
+      auto *val = cast<SetValue>(state.at(set));
+      for (auto *el : val->getAsArrayRef())
+        result.push_back(el);
+
+      if (options.debugMode)
+        for (auto *el : val->getDebugMap())
+          debugMap.push_back(el);
+    }
+
+    if (options.debugMode) {
+      internalizeResult<SetValue>(op.getResult(), std::move(result),
+                                  std::move(debugMap));
+      return DeletionKind::Delete;
+    }
+
+    internalizeResult<SetValue>(op.getResult(), std::move(result));
+    return DeletionKind::Delete;
+  }
+
   FailureOr<DeletionKind>
   visitOp(BagCreateOp op, function_ref<void(Operation *)> addToWorklist) {
     DenseMap<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
@@ -52,3 +52,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
@@ -30,7 +30,9 @@ rtg.test @setOperations : !rtg.dict<> {
   // CHECK-NEXT: }
   %0 = rtg.sequence_closure @seq0
   %1 = rtg.sequence_closure @seq2(%0 : !rtg.sequence)
-  %set = rtg.set_create %0, %1 : !rtg.sequence
+  %set0 = rtg.set_create %0 : !rtg.sequence
+  %set1 = rtg.set_create %1 : !rtg.sequence
+  %set = rtg.set_union %set0, %set1 : !rtg.set<!rtg.sequence>
   %seq = rtg.set_select_random %set : !rtg.set<!rtg.sequence> {rtg.elaboration = 0}
   %new_set = rtg.set_create %seq : !rtg.sequence
   %diff = rtg.set_difference %set, %new_set : !rtg.set<!rtg.sequence>