Skip to content

Commit

Permalink
[RTG] Elaboration support for set_size and bag_unique_size operations (
Browse files Browse the repository at this point in the history
  • Loading branch information
maerhart authored Dec 9, 2024
1 parent 7f8c80b commit d2a3117
Show file tree
Hide file tree
Showing 5 changed files with 49 additions and 4 deletions.
10 changes: 6 additions & 4 deletions include/circt/Dialect/RTG/IR/RTGVisitors.h
Original file line number Diff line number Diff line change
Expand Up @@ -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...);
Expand Down Expand Up @@ -158,6 +158,8 @@ class RTGTypeVisitor {
HANDLE(SetType, Unhandled);
HANDLE(BagType, Unhandled);
HANDLE(DictType, Unhandled);
HANDLE(IndexType, Unhandled);
HANDLE(IntegerType, Unhandled);
#undef HANDLE
};

Expand Down
16 changes: 16 additions & 0 deletions lib/Dialect/RTG/Transforms/ElaborationPass.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
24 changes: 24 additions & 0 deletions test/Dialect/RTG/Transform/elaboration.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]]) :
Expand Down
1 change: 1 addition & 0 deletions tools/circt-opt/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ target_link_libraries(circt-opt
MLIREmitCDialect
MLIRFuncInlinerExtension
MLIRVectorDialect
MLIRIndexDialect
)

export_executable_symbols_for_plugins(circt-opt)
2 changes: 2 additions & 0 deletions tools/circt-opt/circt-opt.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit d2a3117

Please sign in to comment.