From 882ac5b43cb5ed08e749b2ff6c9597ee7447344c Mon Sep 17 00:00:00 2001 From: tancheng Date: Sun, 15 Jun 2025 07:33:05 +0000 Subject: [PATCH] Avoid inserting unnecessary phi when there is only one predecessor block --- .../TransformCtrlToDataFlowPass.cpp | 33 ++++++++++ test/neura/ctrl/branch.mlir | 65 +++++++++++++++++++ test/neura/ctrl/branch_no_arg.mlir | 12 ++-- 3 files changed, 102 insertions(+), 8 deletions(-) create mode 100644 test/neura/ctrl/branch.mlir diff --git a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp index 0db4c778..ab0335e6 100644 --- a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp +++ b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp @@ -113,6 +113,39 @@ void createPhiNodesForBlock(Block *block, OpBuilder &builder, assert(!phi_operands.empty()); + // Puts all operands into a set to ensure uniqueness. Specifically, following + // case is handled: + // --------------------------------------------------------- + // ^bb1: + // "neura.br"(%a)[^bb3] : (!neura.data) -> () + // + // ^bb2: + // "neura.br"(%a)[^bb3] : (!neura.data) -> () + // + // ^bb3(%x: !neura.data): + // ... + // --------------------------------------------------------- + // In above case, %a is used in both branches of the control flow, so we + // don't need a phi node, but we still need to replace its uses with the + // result of the phi node. + // This ensures that we only create a phi node if there are multiple unique + // operands. + llvm::SmallDenseSet unique_operands(phi_operands.begin(), phi_operands.end()); + + if (unique_operands.size() == 1) { + // No phi needed, but still replace + Value single = *unique_operands.begin(); + SmallVector uses; + for (OpOperand &use : live_in.getUses()) { + uses.push_back(&use); + } + for (OpOperand *use : uses) { + use->set(single); + } + value_map[live_in] = single; + continue; + } + // Creates the phi node with dynamic number of operands. auto phi_op = builder.create(loc, predicated_type, phi_operands); diff --git a/test/neura/ctrl/branch.mlir b/test/neura/ctrl/branch.mlir new file mode 100644 index 00000000..606ff758 --- /dev/null +++ b/test/neura/ctrl/branch.mlir @@ -0,0 +1,65 @@ +// RUN: mlir-neura-opt %s \ +// RUN: --assign-accelerator \ +// RUN: --lower-llvm-to-neura \ +// RUN: --leverage-predicated-value \ +// RUN: | FileCheck %s + +// RUN: mlir-neura-opt %s \ +// RUN: --assign-accelerator \ +// RUN: --lower-llvm-to-neura \ +// RUN: --leverage-predicated-value \ +// RUN: --transform-ctrl-to-data-flow \ +// RUN: | FileCheck %s -check-prefix=CTRL2DATA + +func.func @test(%in: i64) -> f32 { + %c0 = llvm.mlir.constant(0 : i64) : i64 + %c1 = llvm.mlir.constant(1.0 : f32) : f32 + %c2 = llvm.mlir.constant(2.0 : f32) : f32 + %c3 = llvm.mlir.constant(3.0 : f32) : f32 + %c4 = llvm.mlir.constant(4.0 : f32) : f32 + %cond = llvm.icmp "eq" %in, %c0 : i64 + llvm.cond_br %cond, ^bb2(%c3, %c4 : f32, f32), ^bb1(%c1, %c2 : f32, f32) + +^bb1(%ca: f32, %cb: f32): + %a = llvm.fadd %ca, %cb : f32 + llvm.br ^bb3(%a : f32) + +^bb2(%cc: f32, %cd: f32): + %b = llvm.fmul %cc, %cd : f32 + llvm.br ^bb3(%b : f32) + +^bb3(%v: f32): + return %v : f32 +} + + +// CHECK: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} { +// CHECK-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data +// CHECK-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data +// CHECK-NEXT: %2 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data +// CHECK-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data +// CHECK-NEXT: %4 = "neura.constant"() <{predicate = true, value = 4.000000e+00 : f32}> : () -> !neura.data +// CHECK-NEXT: %5 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data) -> !neura.data +// CHECK-NEXT: neura.cond_br %5 : !neura.data then %3, %4 : !neura.data, !neura.data to ^bb2 else %1, %2 : !neura.data, !neura.data to ^bb1 +// CHECK-NEXT: ^bb1(%6: !neura.data, %7: !neura.data): // pred: ^bb0 +// CHECK-NEXT: %8 = "neura.fadd"(%6, %7) : (!neura.data, !neura.data) -> !neura.data +// CHECK-NEXT: neura.br %8 : !neura.data to ^bb3 +// CHECK-NEXT: ^bb2(%9: !neura.data, %10: !neura.data): // pred: ^bb0 +// CHECK-NEXT: %11 = "neura.fmul"(%9, %10) : (!neura.data, !neura.data) -> !neura.data +// CHECK-NEXT: neura.br %11 : !neura.data to ^bb3 +// CHECK-NEXT: ^bb3(%12: !neura.data): // 2 preds: ^bb1, ^bb2 +// CHECK-NEXT: "neura.return"(%12) : (!neura.data) -> () +// CHECK-NEXT: } + +// CTRL2DATA: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} { +// CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data +// CTRL2DATA-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data +// CTRL2DATA-NEXT: %2 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data +// CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data +// CTRL2DATA-NEXT: %4 = "neura.constant"() <{predicate = true, value = 4.000000e+00 : f32}> : () -> !neura.data +// CTRL2DATA-NEXT: %5 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %6 = "neura.fadd"(%1, %2) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %7 = "neura.fmul"(%3, %4) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %8 = "neura.phi"(%6, %7) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: "neura.return"(%8) : (!neura.data) -> () +// CTRL2DATA-NEXT: } \ No newline at end of file diff --git a/test/neura/ctrl/branch_no_arg.mlir b/test/neura/ctrl/branch_no_arg.mlir index 26165db5..4755fd1c 100644 --- a/test/neura/ctrl/branch_no_arg.mlir +++ b/test/neura/ctrl/branch_no_arg.mlir @@ -57,12 +57,8 @@ func.func @test(%in: i64) -> f32 { // CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data // CTRL2DATA-NEXT: %4 = "neura.constant"() <{predicate = true, value = 4.000000e+00 : f32}> : () -> !neura.data // CTRL2DATA-NEXT: %5 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data) -> !neura.data -// CTRL2DATA-NEXT: %6 = "neura.phi"(%1) : (!neura.data) -> !neura.data -// CTRL2DATA-NEXT: %7 = "neura.phi"(%2) : (!neura.data) -> !neura.data -// CTRL2DATA-NEXT: %8 = "neura.fadd"(%6, %7) : (!neura.data, !neura.data) -> !neura.data -// CTRL2DATA-NEXT: %9 = "neura.phi"(%3) : (!neura.data) -> !neura.data -// CTRL2DATA-NEXT: %10 = "neura.phi"(%4) : (!neura.data) -> !neura.data -// CTRL2DATA-NEXT: %11 = "neura.fmul"(%9, %10) : (!neura.data, !neura.data) -> !neura.data -// CTRL2DATA-NEXT: %12 = "neura.phi"(%8, %11) : (!neura.data, !neura.data) -> !neura.data -// CTRL2DATA-NEXT: "neura.return"(%12) : (!neura.data) -> () +// CTRL2DATA-NEXT: %6 = "neura.fadd"(%1, %2) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %7 = "neura.fmul"(%3, %4) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %8 = "neura.phi"(%6, %7) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: "neura.return"(%8) : (!neura.data) -> () // CTRL2DATA-NEXT: }