Skip to content

Enale deterministic ctrol to data flow transform #65

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ShangkunLi merged 3 commits into from
Jul 2, 2025
Merged

Enale deterministic ctrol to data flow transform #65

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
14bc44f
enbale deterministic ctrl2data transform
ShangkunLi Jul 1, 2025
e9a64ce
[fix] remove the issue link
ShangkunLi Jul 1, 2025
5e509f6
add comments
ShangkunLi Jul 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 22 additions & 26 deletions lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/raw_ostream.h"
Expand All @@ -25,9 +26,6 @@ using namespace mlir;
#define GEN_PASS_DEF_TransformCtrlToDataFlow
#include "NeuraDialect/NeuraPasses.h.inc"

// TODO: Needs to enbale a deterministic ctrol to data flow transformation
// https://github.com/coredac/dataflow/issues/64

// Inserts `grant_once` for every predicated value defined in the entry block
// that is used outside of the block (i.e., a live-out).
void GrantPredicateInEntryBlock(Block *entry_block, OpBuilder &builder) {
Expand Down Expand Up @@ -115,15 +113,16 @@ struct ControlFlowInfo {
SmallVector<Value> passed_values; // Values passed to the target block.
bool is_back_edge;
};
std::vector<std::unique_ptr<Edge>> all_edges; // All edges in the function.
DenseMap<Block *, SmallVector<Edge *>>
SmallVector<std::unique_ptr<Edge>> all_edges; // All edges in the function.
llvm::MapVector<Block *, SmallVector<Edge *>>
incoming_edges; // Incoming edges for each block.
DenseMap<Block *, SmallVector<Edge *>>
llvm::MapVector<Block *, SmallVector<Edge *>>
outgoing_edges; // Outgoing edges for each block.

DenseMap<Block *, SmallVector<Edge *>> back_edges;
DenseMap<Block *, SmallVector<Edge *>> forward_edges;
DenseSet<Block *> blocks_with_back_edges; // Blocks with backward edges.
llvm::MapVector<Block *, SmallVector<Edge *>> back_edges;
llvm::MapVector<Block *, SmallVector<Edge *>> forward_edges;
llvm::SmallVector<Block *>
blocks_with_back_edges; // Blocks with backward edges.

Edge *createEdge() {
all_edges.push_back(std::make_unique<Edge>());
Expand Down Expand Up @@ -235,13 +234,13 @@ void buildControlFlowInfo(func::FuncOp &func, ControlFlowInfo &ctrl_info,
// Handles back edges.
if (true_edge->is_back_edge) {
ctrl_info.back_edges[&block].push_back(true_edge);
ctrl_info.blocks_with_back_edges.insert(&block);
ctrl_info.blocks_with_back_edges.push_back(&block);
} else {
ctrl_info.forward_edges[&block].push_back(true_edge);
}
if (false_edge->is_back_edge) {
ctrl_info.back_edges[&block].push_back(false_edge);
ctrl_info.blocks_with_back_edges.insert(&block);
ctrl_info.blocks_with_back_edges.push_back(&block);
} else {
ctrl_info.forward_edges[&block].push_back(false_edge);
}
Expand All @@ -267,7 +266,7 @@ void buildControlFlowInfo(func::FuncOp &func, ControlFlowInfo &ctrl_info,
// Handles back edges.
if (edge->is_back_edge) {
ctrl_info.back_edges[&block].push_back(edge);
ctrl_info.blocks_with_back_edges.insert(&block);
ctrl_info.blocks_with_back_edges.push_back(&block);
} else {
ctrl_info.forward_edges[&block].push_back(edge);
}
Expand All @@ -282,7 +281,7 @@ void buildControlFlowInfo(func::FuncOp &func, ControlFlowInfo &ctrl_info,
}

Value getPrecessedCondition(Value condition, bool is_not_condition,
DenseMap<Value, Value> &condition_cache,
llvm::MapVector<Value, Value> &condition_cache,
OpBuilder &builder) {
if (!is_not_condition) {
return condition;
Expand All @@ -302,8 +301,8 @@ Value getPrecessedCondition(Value condition, bool is_not_condition,
}

void createReserveAndPhiOps(func::FuncOp &func, ControlFlowInfo &ctrl_info,
DenseMap<BlockArgument, Value> &arg_to_reserve,
DenseMap<BlockArgument, Value> &arg_to_phi_result,
llvm::MapVector<BlockArgument, Value> &arg_to_reserve,
llvm::MapVector<BlockArgument, Value> &arg_to_phi_result,
OpBuilder &builder) {
DominanceInfo dom_info(func);

Expand All @@ -317,17 +316,16 @@ void createReserveAndPhiOps(func::FuncOp &func, ControlFlowInfo &ctrl_info,
// Type 5: Forward cond_br edges without values.
// Type 6: Forward br edges without values.
// For Backward edges without values, they can be transformed into type 1 or 2

DenseMap<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
llvm::MapVector<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
backward_value_edges;
DenseMap<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
llvm::MapVector<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
forward_value_edges;
DenseMap<Block *, SmallVector<ControlFlowInfo::Edge *>>
llvm::MapVector<Block *, SmallVector<ControlFlowInfo::Edge *>>
block_conditional_edges;

DenseMap<Value, Value> condition_cache;
llvm::MapVector<Value, Value> condition_cache;

DenseMap<BlockArgument, SmallVector<Value>> arg_to_phi_operands;
llvm::MapVector<BlockArgument, SmallVector<Value>> arg_to_phi_operands;

for (auto &edge : ctrl_info.all_edges) {
Block *target = edge->target;
Expand Down Expand Up @@ -431,8 +429,6 @@ void createReserveAndPhiOps(func::FuncOp &func, ControlFlowInfo &ctrl_info,
// ================================================
// Step 4: Creates phi operations for each block argument.
// ================================================
DenseSet<BlockArgument> args_needing_phi;

for (auto &arg_to_phi_pair : arg_to_phi_operands) {
BlockArgument arg = arg_to_phi_pair.first;
auto &phi_operands = arg_to_phi_pair.second;
Expand Down Expand Up @@ -501,7 +497,7 @@ void createReserveAndPhiOps(func::FuncOp &func, ControlFlowInfo &ctrl_info,

if (target->getArguments().empty()) {
// Grants predicate for all the live-in values in the target block.
DenseSet<Value> live_in_values;
SetVector<Value> live_in_values;
for (Operation &op : target->getOperations()) {
for (Value operand : op.getOperands()) {
if (operand.getDefiningOp() &&
Expand Down Expand Up @@ -562,8 +558,8 @@ void transformControlFlowToDataFlow(func::FuncOp &func,
assertLiveOutValuesDominatedByBlockArgs(func);

// Creates reserve and phi operations for each block argument.
DenseMap<BlockArgument, Value> arg_to_reserve;
DenseMap<BlockArgument, Value> arg_to_phi_result;
llvm::MapVector<BlockArgument, Value> arg_to_reserve;
llvm::MapVector<BlockArgument, Value> arg_to_phi_result;
createReserveAndPhiOps(func, ctrl_info, arg_to_reserve, arg_to_phi_result,
builder);

Expand Down
18 changes: 9 additions & 9 deletions test/affine2neura/bert/bert_node1/bert_node1.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,19 +66,19 @@ module attributes {} {
// CTRL2DATA-NEXT: %17 = "neura.cast"(%16) <{cast_type = "int_to_index"}> : (!neura.data<i64, i1>) -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %18 = "neura.icmp"(%17, %3) <{cmpType = "slt"}> : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<i1, i1>
// CTRL2DATA-NEXT: %19 = "neura.not"(%18) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
// CTRL2DATA-DAG: %[[VAL1:.*]] = neura.grant_predicate %5, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-DAG: %[[VAL2:.*]] = neura.grant_predicate %17, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %22 = neura.load_indexed %arg0[%[[VAL1:.*]], %[[VAL1:.*]], %[[VAL1:.*]], %[[VAL1:.*]], %[[VAL1:.*]], %[[VAL2:.*]] : !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>] memref<?x1x1x1x1x128xi8> : !neura.data<i8, i1>
// CTRL2DATA-NEXT: %20 = neura.grant_predicate %5, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %21 = neura.grant_predicate %17, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %22 = neura.load_indexed %arg0[%20, %20, %20, %20, %20, %21 : !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>] memref<?x1x1x1x1x128xi8> : !neura.data<i8, i1>
// CTRL2DATA-NEXT: %23 = neura.grant_predicate %10, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: neura.store_indexed %22 to %arg1[%[[VAL1:.*]], %[[VAL1:.*]], %23, %[[VAL1:.*]], %[[VAL1:.*]], %[[VAL2:.*]] : !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>] memref<?x1x128x1x1x128xi8> : !neura.data<i8, i1>
// CTRL2DATA-NEXT: neura.store_indexed %22 to %arg1[%20, %20, %23, %20, %20, %21 : !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>, !neura.data<index, i1>] memref<?x1x128x1x1x128xi8> : !neura.data<i8, i1>
// CTRL2DATA-NEXT: %24 = neura.grant_predicate %1, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %25 = "neura.add"(%[[VAL2:.*]], %24) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %25 = "neura.add"(%21, %24) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %26 = "neura.cast"(%25) <{cast_type = "index_to_int"}> : (!neura.data<index, i1>) -> !neura.data<i64, i1>
// CTRL2DATA-NEXT: neura.ctrl_mov %26 -> %15 : !neura.data<i64, i1> !neura.data<i64, i1>
// CTRL2DATA-DAG: %[[VAL3:.*]] = neura.grant_predicate %10, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-DAG: %[[VAL4:.*]] = neura.grant_predicate %1, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %29 = "neura.add"(%[[VAL3:.*]], %[[VAL4:.*]]) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %27 = neura.grant_predicate %10, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %28 = neura.grant_predicate %1, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %29 = "neura.add"(%27, %28) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
// CTRL2DATA-NEXT: %30 = "neura.cast"(%29) <{cast_type = "index_to_int"}> : (!neura.data<index, i1>) -> !neura.data<i64, i1>
// CTRL2DATA-NEXT: neura.ctrl_mov %30 -> %8 : !neura.data<i64, i1> !neura.data<i64, i1>
// CTRL2DATA-NEXT: "neura.return"() : () -> ()
// CTRL2DATA-NEXT: }
// CTRL2DATA-NEXT: }
Loading