[RTG] Elaboration support for set_size and bag_unique_size operations

llvm · Dec 9, 2024 · 81f191f · 81f191f
1 parent 7f8c80b
commit 81f191f
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 4 deletions.
diff --git a/include/circt/Dialect/RTG/IR/RTGVisitors.h b/include/circt/Dialect/RTG/IR/RTGVisitors.h
@@ -111,10 +111,10 @@ class RTGTypeVisitor {
   ResultType dispatchTypeVisitor(Type type, ExtraArgs... args) {
     auto *thisCast = static_cast<ConcreteType *>(this);
     return TypeSwitch<Type, ResultType>(type)
-        .template Case<SequenceType, SetType, BagType, DictType>(
-            [&](auto expr) -> ResultType {
-              return thisCast->visitType(expr, args...);
-            })
+        .template Case<SequenceType, SetType, BagType, DictType, IndexType,
+                       IntegerType>([&](auto expr) -> ResultType {
+          return thisCast->visitType(expr, args...);
+        })
         .template Case<ContextResourceTypeInterface>(
             [&](auto expr) -> ResultType {
               return thisCast->visitContextResourceType(expr, args...);
@@ -158,6 +158,8 @@ class RTGTypeVisitor {
   HANDLE(SetType, Unhandled);
   HANDLE(BagType, Unhandled);
   HANDLE(DictType, Unhandled);
+  HANDLE(IndexType, Unhandled);
+  HANDLE(IntegerType, Unhandled);
 #undef HANDLE
 };
 

diff --git a/lib/Dialect/RTG/Transforms/ElaborationPass.cpp b/lib/Dialect/RTG/Transforms/ElaborationPass.cpp
@@ -525,6 +525,14 @@ class Elaborator : public RTGOpVisitor<Elaborator, FailureOr<DeletionKind>,
     return DeletionKind::Delete;
   }
 
+  FailureOr<DeletionKind>
+  visitOp(SetSizeOp op, function_ref<void(Operation *)> addToWorklist) {
+    auto size = cast<SetValue>(state.at(op.getSet()))->getSet().size();
+    auto sizeAttr = IntegerAttr::get(IndexType::get(op->getContext()), size);
+    internalizeResult<AttributeValue>(op.getResult(), sizeAttr);
+    return DeletionKind::Delete;
+  }
+
   FailureOr<DeletionKind>
   visitOp(BagCreateOp op, function_ref<void(Operation *)> addToWorklist) {
     MapVector<ElaboratorValue *, uint64_t> bag;
@@ -612,6 +620,14 @@ class Elaborator : public RTGOpVisitor<Elaborator, FailureOr<DeletionKind>,
     return DeletionKind::Delete;
   }
 
+  FailureOr<DeletionKind>
+  visitOp(BagUniqueSizeOp op, function_ref<void(Operation *)> addToWorklist) {
+    auto size = cast<BagValue>(state.at(op.getBag()))->getBag().size();
+    auto sizeAttr = IntegerAttr::get(IndexType::get(op->getContext()), size);
+    internalizeResult<AttributeValue>(op.getResult(), sizeAttr);
+    return DeletionKind::Delete;
+  }
+
   FailureOr<DeletionKind>
   dispatchOpVisitor(Operation *op,
                     function_ref<void(Operation *)> addToWorklist) {

diff --git a/test/Dialect/RTG/Transform/elaboration.mlir b/test/Dialect/RTG/Transform/elaboration.mlir
@@ -4,6 +4,7 @@ func.func @dummy1(%arg0: i32, %arg1: i32, %arg2: !rtg.set<i32>) -> () {return}
 func.func @dummy2(%arg0: i32) -> () {return}
 func.func @dummy3(%arg0: i64) -> () {return}
 func.func @dummy4(%arg0: i32, %arg1: i32, %arg2: !rtg.bag<i32>, %arg3: !rtg.bag<i32>) -> () {return}
+func.func @dummy5(%arg0: index) -> () {return}
 
 // Test the set operations and passing a sequence to another one via argument
 // CHECK-LABEL: rtg.test @setOperations
@@ -54,6 +55,29 @@ rtg.test @bagOperations : !rtg.dict<> {
   func.call @dummy4(%3, %4, %diff, %diff2) : (i32, i32, !rtg.bag<i32>, !rtg.bag<i32>) -> ()
 }
 
+// CHECK-LABEL: rtg.test @setSize
+rtg.test @setSize : !rtg.dict<> {
+  // CHECK-NEXT: [[C:%.+]] = arith.constant 1 : index
+  // CHECK-NEXT: func.call @dummy5([[C]])
+  // CHECK-NEXT: }
+  %c5_i32 = arith.constant 5 : i32
+  %set = rtg.set_create %c5_i32 : i32
+  %size = rtg.set_size %set : !rtg.set<i32>
+  func.call @dummy5(%size) : (index) -> ()
+}
+
+// CHECK-LABEL: rtg.test @bagSize
+rtg.test @bagSize : !rtg.dict<> {
+  // CHECK-NEXT: [[C:%.+]] = arith.constant 1 : index
+  // CHECK-NEXT: func.call @dummy5([[C]])
+  // CHECK-NEXT: }
+  %c8 = arith.constant 8 : index
+  %c5_i32 = arith.constant 5 : i32
+  %bag = rtg.bag_create (%c8 x %c5_i32) : i32
+  %size = rtg.bag_unique_size %bag : !rtg.bag<i32>
+  func.call @dummy5(%size) : (index) -> ()
+}
+
 // CHECK-LABEL: @targetTest_target0
 // CHECK: [[V0:%.+]] = arith.constant 0
 // CHECK: func.call @dummy2([[V0]]) :

diff --git a/tools/circt-opt/CMakeLists.txt b/tools/circt-opt/CMakeLists.txt
@@ -38,6 +38,7 @@ target_link_libraries(circt-opt
   MLIREmitCDialect
   MLIRFuncInlinerExtension
   MLIRVectorDialect
+  MLIRIndexDialect
 )
 
 export_executable_symbols_for_plugins(circt-opt)
diff --git a/tools/circt-opt/circt-opt.cpp b/tools/circt-opt/circt-opt.cpp
@@ -22,6 +22,7 @@
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -55,6 +56,7 @@ int main(int argc, char **argv) {
   registry.insert<mlir::scf::SCFDialect>();
   registry.insert<mlir::emitc::EmitCDialect>();
   registry.insert<mlir::vector::VectorDialect>();
+  registry.insert<mlir::index::IndexDialect>();
 
   circt::registerAllDialects(registry);
   circt::registerAllPasses();