Skip to content

Commit 05f3a43

Browse files
ShangkunLiShangkunLi
authored
Merge pull request #65 from ShangkunLi/determine-ctrl2data
Enale deterministic ctrol to data flow transform
2 parents 394efff + db2aaeb commit 05f3a43

File tree

7 files changed

+162
-163
lines changed

7 files changed

+162
-163
lines changed

‎lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp‎

Lines changed: 24 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include "mlir/Pass/Pass.h"
1515
#include "mlir/Support/LLVM.h"
1616
#include "llvm/ADT/STLExtras.h"
17+
#include "llvm/ADT/MapVector.h"
1718
#include "llvm/Support/Casting.h"
1819
#include "llvm/Support/Error.h"
1920
#include "llvm/Support/raw_ostream.h"
@@ -25,9 +26,6 @@ using namespace mlir;
2526
#define GEN_PASS_DEF_TransformCtrlToDataFlow
2627
#include "NeuraDialect/NeuraPasses.h.inc"
2728

28-
// TODO: Needs to enbale a deterministic ctrol to data flow transformation
29-
// https://github.com/coredac/dataflow/issues/64
30-
3129
// Inserts `grant_once` for every predicated value defined in the entry block
3230
// that is used outside of the block (i.e., a live-out).
3331
void GrantPredicateInEntryBlock(Block *entry_block, OpBuilder &builder) {
@@ -115,15 +113,16 @@ struct ControlFlowInfo {
115113
SmallVector<Value> passed_values; // Values passed to the target block.
116114
bool is_back_edge;
117115
};
118-
std::vector<std::unique_ptr<Edge>> all_edges; // All edges in the function.
119-
DenseMap<Block *, SmallVector<Edge *>>
116+
SmallVector<std::unique_ptr<Edge>> all_edges; // All edges in the function.
117+
llvm::MapVector<Block *, SmallVector<Edge *>>
120118
incoming_edges; // Incoming edges for each block.
121-
DenseMap<Block *, SmallVector<Edge *>>
119+
llvm::MapVector<Block *, SmallVector<Edge *>>
122120
outgoing_edges; // Outgoing edges for each block.
123121

124-
DenseMap<Block *, SmallVector<Edge *>> back_edges;
125-
DenseMap<Block *, SmallVector<Edge *>> forward_edges;
126-
DenseSet<Block *> blocks_with_back_edges; // Blocks with backward edges.
122+
llvm::MapVector<Block *, SmallVector<Edge *>> back_edges;
123+
llvm::MapVector<Block *, SmallVector<Edge *>> forward_edges;
124+
llvm::SmallVector<Block *>
125+
blocks_with_back_edges; // Blocks with backward edges.
127126

128127
Edge *createEdge() {
129128
all_edges.push_back(std::make_unique<Edge>());
@@ -235,13 +234,13 @@ void buildControlFlowInfo(func::FuncOp &func, ControlFlowInfo &ctrl_info,
235234
// Handles back edges.
236235
if (true_edge->is_back_edge) {
237236
ctrl_info.back_edges[&block].push_back(true_edge);
238-
ctrl_info.blocks_with_back_edges.insert(&block);
237+
ctrl_info.blocks_with_back_edges.push_back(&block);
239238
} else {
240239
ctrl_info.forward_edges[&block].push_back(true_edge);
241240
}
242241
if (false_edge->is_back_edge) {
243242
ctrl_info.back_edges[&block].push_back(false_edge);
244-
ctrl_info.blocks_with_back_edges.insert(&block);
243+
ctrl_info.blocks_with_back_edges.push_back(&block);
245244
} else {
246245
ctrl_info.forward_edges[&block].push_back(false_edge);
247246
}
@@ -267,7 +266,7 @@ void buildControlFlowInfo(func::FuncOp &func, ControlFlowInfo &ctrl_info,
267266
// Handles back edges.
268267
if (edge->is_back_edge) {
269268
ctrl_info.back_edges[&block].push_back(edge);
270-
ctrl_info.blocks_with_back_edges.insert(&block);
269+
ctrl_info.blocks_with_back_edges.push_back(&block);
271270
} else {
272271
ctrl_info.forward_edges[&block].push_back(edge);
273272
}
@@ -282,7 +281,7 @@ void buildControlFlowInfo(func::FuncOp &func, ControlFlowInfo &ctrl_info,
282281
}
283282

284283
Value getPrecessedCondition(Value condition, bool is_not_condition,
285-
DenseMap<Value, Value> &condition_cache,
284+
llvm::MapVector<Value, Value> &condition_cache,
286285
OpBuilder &builder) {
287286
if (!is_not_condition) {
288287
return condition;
@@ -302,8 +301,8 @@ Value getPrecessedCondition(Value condition, bool is_not_condition,
302301
}
303302

304303
void createReserveAndPhiOps(func::FuncOp &func, ControlFlowInfo &ctrl_info,
305-
DenseMap<BlockArgument, Value> &arg_to_reserve,
306-
DenseMap<BlockArgument, Value> &arg_to_phi_result,
304+
llvm::MapVector<BlockArgument, Value> &arg_to_reserve,
305+
llvm::MapVector<BlockArgument, Value> &arg_to_phi_result,
307306
OpBuilder &builder) {
308307
DominanceInfo dom_info(func);
309308

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

321-
DenseMap<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
320+
// Uses llvm::MapVector instead of DenseMap to maintain insertion order.
321+
llvm::MapVector<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
322322
backward_value_edges;
323-
DenseMap<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
323+
llvm::MapVector<BlockArgument, SmallVector<ControlFlowInfo::Edge *>>
324324
forward_value_edges;
325-
DenseMap<Block *, SmallVector<ControlFlowInfo::Edge *>>
325+
llvm::MapVector<Block *, SmallVector<ControlFlowInfo::Edge *>>
326326
block_conditional_edges;
327327

328-
DenseMap<Value, Value> condition_cache;
328+
llvm::MapVector<Value, Value> condition_cache;
329329

330-
DenseMap<BlockArgument, SmallVector<Value>> arg_to_phi_operands;
330+
llvm::MapVector<BlockArgument, SmallVector<Value>> arg_to_phi_operands;
331331

332332
for (auto &edge : ctrl_info.all_edges) {
333333
Block *target = edge->target;
@@ -431,8 +431,6 @@ void createReserveAndPhiOps(func::FuncOp &func, ControlFlowInfo &ctrl_info,
431431
// ================================================
432432
// Step 4: Creates phi operations for each block argument.
433433
// ================================================
434-
DenseSet<BlockArgument> args_needing_phi;
435-
436434
for (auto &arg_to_phi_pair : arg_to_phi_operands) {
437435
BlockArgument arg = arg_to_phi_pair.first;
438436
auto &phi_operands = arg_to_phi_pair.second;
@@ -501,7 +499,8 @@ void createReserveAndPhiOps(func::FuncOp &func, ControlFlowInfo &ctrl_info,
501499

502500
if (target->getArguments().empty()) {
503501
// Grants predicate for all the live-in values in the target block.
504-
DenseSet<Value> live_in_values;
502+
// Uses SetVector instead of DenseSet to maintain insertion order.
503+
SetVector<Value> live_in_values;
505504
for (Operation &op : target->getOperations()) {
506505
for (Value operand : op.getOperands()) {
507506
if (operand.getDefiningOp() &&
@@ -562,8 +561,8 @@ void transformControlFlowToDataFlow(func::FuncOp &func,
562561
assertLiveOutValuesDominatedByBlockArgs(func);
563562

564563
// Creates reserve and phi operations for each block argument.
565-
DenseMap<BlockArgument, Value> arg_to_reserve;
566-
DenseMap<BlockArgument, Value> arg_to_phi_result;
564+
llvm::MapVector<BlockArgument, Value> arg_to_reserve;
565+
llvm::MapVector<BlockArgument, Value> arg_to_phi_result;
567566
createReserveAndPhiOps(func, ctrl_info, arg_to_reserve, arg_to_phi_result,
568567
builder);
569568

‎test/affine2neura/bert/bert_node1/bert_node1.mlir‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -66,19 +66,19 @@ module attributes {} {
6666
// CTRL2DATA-NEXT: %17 = "neura.cast"(%16) <{cast_type = "int_to_index"}> : (!neura.data<i64, i1>) -> !neura.data<index, i1>
6767
// CTRL2DATA-NEXT: %18 = "neura.icmp"(%17, %3) <{cmpType = "slt"}> : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<i1, i1>
6868
// CTRL2DATA-NEXT: %19 = "neura.not"(%18) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
69-
// CTRL2DATA-DAG: %[[VAL1:.*]] = neura.grant_predicate %5, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
70-
// CTRL2DATA-DAG: %[[VAL2:.*]] = neura.grant_predicate %17, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
71-
// 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>
69+
// CTRL2DATA-NEXT: %20 = neura.grant_predicate %5, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
70+
// CTRL2DATA-NEXT: %21 = neura.grant_predicate %17, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
71+
// 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>
7272
// CTRL2DATA-NEXT: %23 = neura.grant_predicate %10, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
73-
// 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>
73+
// 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>
7474
// CTRL2DATA-NEXT: %24 = neura.grant_predicate %1, %18 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
75-
// CTRL2DATA-NEXT: %25 = "neura.add"(%[[VAL2:.*]], %24) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
75+
// CTRL2DATA-NEXT: %25 = "neura.add"(%21, %24) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
7676
// CTRL2DATA-NEXT: %26 = "neura.cast"(%25) <{cast_type = "index_to_int"}> : (!neura.data<index, i1>) -> !neura.data<i64, i1>
7777
// CTRL2DATA-NEXT: neura.ctrl_mov %26 -> %15 : !neura.data<i64, i1> !neura.data<i64, i1>
78-
// CTRL2DATA-DAG: %[[VAL3:.*]] = neura.grant_predicate %10, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
79-
// CTRL2DATA-DAG: %[[VAL4:.*]] = neura.grant_predicate %1, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
80-
// CTRL2DATA-NEXT: %29 = "neura.add"(%[[VAL3:.*]], %[[VAL4:.*]]) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
78+
// CTRL2DATA-NEXT: %27 = neura.grant_predicate %10, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
79+
// CTRL2DATA-NEXT: %28 = neura.grant_predicate %1, %19 : !neura.data<index, i1>, !neura.data<i1, i1> -> !neura.data<index, i1>
80+
// CTRL2DATA-NEXT: %29 = "neura.add"(%27, %28) : (!neura.data<index, i1>, !neura.data<index, i1>) -> !neura.data<index, i1>
8181
// CTRL2DATA-NEXT: %30 = "neura.cast"(%29) <{cast_type = "index_to_int"}> : (!neura.data<index, i1>) -> !neura.data<i64, i1>
8282
// CTRL2DATA-NEXT: neura.ctrl_mov %30 -> %8 : !neura.data<i64, i1> !neura.data<i64, i1>
8383
// CTRL2DATA-NEXT: "neura.return"() : () -> ()
84-
// CTRL2DATA-NEXT: }
84+
// CTRL2DATA-NEXT: }

0 commit comments

Comments
 (0)