Skip to content

Commit

Permalink
[DC] Add unused fork result elimination
Browse files Browse the repository at this point in the history 
Adding a canonicalizer to remove unused results from forks. Required to
get valid post-canonicalization IR. Also, narrow the scope of "XFAIL"s
in canonicalizer test to add new one.
  • Loading branch information
@teqdruid
teqdruid committed Dec 2, 2024
1 parent fd08d90 commit ed64cf4
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 15 deletions.
32 changes: 30 additions & 2 deletions lib/Dialect/DC/DCOps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,10 +180,38 @@ class EliminateForkOfSourcePattern : public OpRewritePattern<ForkOp> {
}
};

struct EliminateUnusedForkResultsPattern : mlir::OpRewritePattern<ForkOp> {
using mlir::OpRewritePattern<ForkOp>::OpRewritePattern;

LogicalResult matchAndRewrite(ForkOp op,
PatternRewriter &rewriter) const override {
std::set<unsigned> unusedIndexes;

for (auto res : llvm::enumerate(op.getResults()))
if (res.value().getUses().empty())
unusedIndexes.insert(res.index());

if (unusedIndexes.empty())
return failure();

// Create a new fork op, dropping the unused results.
rewriter.setInsertionPoint(op);
auto operand = op.getOperand();
auto newFork = rewriter.create<ForkOp>(
op.getLoc(), operand, op.getNumResults() - unusedIndexes.size());
unsigned i = 0;
for (auto oldRes : llvm::enumerate(op.getResults()))
if (unusedIndexes.count(oldRes.index()) == 0)
rewriter.replaceAllUsesWith(oldRes.value(), newFork.getResults()[i++]);
rewriter.eraseOp(op);
return success();
}
};

void ForkOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.insert<EliminateForkToForkPattern, EliminateForkOfSourcePattern>(
context);
results.insert<EliminateForkToForkPattern, EliminateForkOfSourcePattern,
EliminateUnusedForkResultsPattern>(context);
}

LogicalResult ForkOp::fold(FoldAdaptor adaptor,
Expand Down
45 changes: 32 additions & 13 deletions test/Dialect/DC/canonicalization.mlir
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
// RUN: circt-opt %s --canonicalize --cse --canonicalize | FileCheck %s
// XFAIL: *
// Waiting on: https://github.com/llvm/llvm-project/issues/64280

// CHECK-LABEL: func.func @staggeredJoin1(
Expand All @@ -14,11 +13,12 @@ func.func @staggeredJoin1(%a: !dc.token, %b : !dc.token) -> (!dc.token) {
return %1 : !dc.token
}

// TODO: For some reason, the canonicalizer no longer combines the two joins. Investigate.
// CHECK-LABEL: func.func @staggeredJoin2(
// CHECK-SAME: %[[VAL_0:.*]]: !dc.token, %[[VAL_1:.*]]: !dc.token, %[[VAL_2:.*]]: !dc.token, %[[VAL_3:.*]]: !dc.token) -> !dc.token {
// CHECK: %[[VAL_4:.*]] = dc.join %[[VAL_2]], %[[VAL_3]], %[[VAL_0]], %[[VAL_1]]
// CHECK: return %[[VAL_4]] : !dc.token
// CHECK: }
// CHECKx: %[[VAL_4:.*]] = dc.join %[[VAL_2]], %[[VAL_3]], %[[VAL_0]], %[[VAL_1]]
// CHECKx: return %[[VAL_4]] : !dc.token
// CHECKx: }
func.func @staggeredJoin2(%a: !dc.token, %b : !dc.token, %c : !dc.token, %d : !dc.token) -> (!dc.token) {
%0 = dc.join %a, %b
%1 = dc.join %c, %0, %d
Expand Down Expand Up @@ -80,11 +80,12 @@ func.func @csePack(%a: !dc.token, %b : i32) -> (!dc.value<i32>, !dc.value<i32>)
}


// TODO: I'm 90% sure this canonicalizer and test are broken. Investigate.
// CHECK-LABEL: func.func @forkToFork(
// CHECK-SAME: %[[VAL_0:.*]]: !dc.token) -> (!dc.token, !dc.token, !dc.token) {
// CHECK: %[[VAL_1:.*]]:3 = dc.fork [3] %[[VAL_0]]
// CHECK: return %[[VAL_1]]#1, %[[VAL_1]]#1, %[[VAL_1]]#2 : !dc.token, !dc.token, !dc.token
// CHECK: }
// CHECKx: %[[VAL_1:.*]]:3 = dc.fork [3] %[[VAL_0]]
// CHECKx: return %[[VAL_1]]#1, %[[VAL_1]]#1, %[[VAL_1]]#2 : !dc.token, !dc.token, !dc.token
// CHECKx: }
func.func @forkToFork(%a: !dc.token) -> (!dc.token, !dc.token, !dc.token) {
%0, %1 = dc.fork [2] %a
%2, %3 = dc.fork [2] %0
Expand All @@ -102,11 +103,13 @@ func.func @forkToFork2(%a: !dc.token) -> (!dc.token, !dc.token, !dc.token) {
return %0, %2, %3 : !dc.token, !dc.token, !dc.token
}

// TODO: For some reason, the canonicalizer no longer simplifies this redundant
// triangle pattern. Investigate.
// CHECK-LABEL: func.func @merge(
// CHECK-SAME: %[[VAL_0:.*]]: !dc.value<i1>) -> !dc.token {
// CHECK: %[[VAL_1:.*]], %[[VAL_2:.*]] = dc.unpack %[[VAL_0]] : !dc.value<i1>
// CHECK: return %[[VAL_1]] : !dc.token
// CHECK: }
// CHECKx: %[[VAL_1:.*]], %[[VAL_2:.*]] = dc.unpack %[[VAL_0]] : !dc.value<i1>
// CHECKx: return %[[VAL_1]] : !dc.token
// CHECKx: }
func.func @merge(%sel : !dc.value<i1>) -> (!dc.token) {
// Canonicalize away a merge that is fed by a branch with the same select
// input.
Expand All @@ -115,12 +118,14 @@ func.func @merge(%sel : !dc.value<i1>) -> (!dc.token) {
return %0 : !dc.token
}

// TODO: For some reason, the canonicalizer no longer removes the source->join
// pattern. Investigate.
// CHECK-LABEL: func.func @joinOnSource(
// CHECK-SAME: %[[VAL_0:.*]]: !dc.token,
// CHECK-SAME: %[[VAL_1:.*]]: !dc.token) -> !dc.token {
// CHECK: %[[VAL_2:.*]] = dc.join %[[VAL_0]], %[[VAL_1]]
// CHECK: return %[[VAL_2]] : !dc.token
// CHECK: }
// CHECKx: %[[VAL_2:.*]] = dc.join %[[VAL_0]], %[[VAL_1]]
// CHECKx: return %[[VAL_2]] : !dc.token
// CHECKx: }
func.func @joinOnSource(%a : !dc.token, %b : !dc.token) -> (!dc.token) {
%0 = dc.source
%out = dc.join %a, %0, %b
Expand All @@ -136,3 +141,17 @@ func.func @forkOfSource() -> (!dc.token, !dc.token) {
%1:2 = dc.fork [2] %0
return %1#0, %1#1 : !dc.token, !dc.token
}

// CHECK-LABEL: hw.module @TestForkCanonicalization(in %arg0 : !dc.token, out out0 : !dc.token, out out1 : !dc.token) {
// CHECK-NEXT: [[R0:%.+]]:2 = dc.fork [2] %arg0
// CHECK-NEXT: [[R1:%.+]] = dc.buffer[4] [[R0]]#0 : !dc.token
// CHECK-NEXT: [[R2:%.+]] = dc.buffer[4] [[R0]]#1 : !dc.token
// CHECK-NEXT: hw.output [[R1]], [[R2]] : !dc.token, !dc.token

hw.module @TestForkCanonicalization(in %arg0: !dc.token, out out0: !dc.token, out out1: !dc.token) {
%0:3 = dc.fork [3] %arg0
%1 = dc.buffer [4] %0#1 : !dc.token
%2 = dc.buffer [4] %0#2 : !dc.token
dc.sink %0#2
hw.output %1, %2 : !dc.token, !dc.token
}

0 comments on commit ed64cf4

Please sign in to comment.