Skip to content

Commit 38ca969

Browse files
tanchengtancheng
authored
Merge pull request #50 from coredac/mapping_util
Re-organize mapping util functions
2 parents 495d301 + c6cb1b8 commit 38ca969

File tree

5 files changed

+144
-80
lines changed

5 files changed

+144
-80
lines changed

‎include/NeuraDialect/mapping/mapping_util.h‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#pragma once
2+
3+
#include "mlir/IR/Operation.h"
4+
5+
namespace mlir {
6+
namespace neura {
7+
8+
// Represents a recurrence cycle rooted at a reserve operation and closed by ctrl_mov.
9+
struct RecurrenceCycle {
10+
SmallVector<Operation *> operations; // Ordered list of operations in the cycle.
11+
int length = 0; // Number of operations excluding reserve/ctrl_mov.
12+
};
13+
14+
// Accelerator configuration struct.
15+
struct AcceleratorConfig {
16+
int num_tiles = 4; // Default to 4 tiles if unspecified.
17+
};
18+
19+
// Collects recurrence cycles rooted at reserve and closed by ctrl_mov.
20+
SmallVector<RecurrenceCycle, 4> collectRecurrenceCycles(Operation *func_op);
21+
22+
// Calculates ResMII: ceil(#ops / #tiles).
23+
int calculateResMii(Operation *func_op, const AcceleratorConfig &config);
24+
25+
} // namespace neura
26+
} // namespace mlir

‎lib/NeuraDialect/Transforms/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ add_mlir_library(
99
TransformCtrlToDataFlowPass.cpp
1010
LeveragePredicatedValuePass.cpp
1111
MapToAcceleratorPass.cpp
12+
mapping/mapping_util.cpp
1213

1314
DEPENDS
1415
MLIRNeuraTransformsIncGen

‎lib/NeuraDialect/Transforms/MapToAcceleratorPass.cpp‎

Lines changed: 11 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -4,94 +4,20 @@
44
#include "NeuraDialect/NeuraOps.h"
55
#include "NeuraDialect/NeuraTypes.h"
66
#include "NeuraDialect/NeuraPasses.h"
7+
#include "NeuraDialect/mapping/mapping_util.h"
78
#include "mlir/Dialect/Func/IR/FuncOps.h"
89
#include "mlir/IR/PatternMatch.h"
910
#include "mlir/Pass/Pass.h"
1011
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
1112

1213
using namespace mlir;
14+
using namespace mlir::neura;
1315

1416
#define GEN_PASS_DEF_MapToAccelerator
1517
#include "NeuraDialect/NeuraPasses.h.inc"
1618

1719
namespace {
1820

19-
/// Represents a recurrence cycle rooted at a reserve operation and ending at a ctrl_mov.
20-
/// The cycle consists of a sequence of operations and its corresponding length.
21-
struct RecurrenceCycle {
22-
SmallVector<Operation *> operations; // Ordered list of operations in the cycle.
23-
int length = 0; // Number of operations excluding ctrl_mov and reserve_op.
24-
};
25-
26-
// Traverses (backward) the operation graph starting from the given operation
27-
// towards reserve_value.
28-
void traverseAlongPath(Operation *op, Value reserve_value,
29-
std::deque<Operation *> &current_path,
30-
DenseSet<Operation *> &visited_in_path,
31-
SmallVector<RecurrenceCycle, 4> &collected_paths) {
32-
if (!op || visited_in_path.contains(op))
33-
return;
34-
35-
visited_in_path.insert(op);
36-
current_path.push_front(op);
37-
38-
for (Value operand : op->getOperands()) {
39-
if (operand == reserve_value) {
40-
Operation *res_op = reserve_value.getDefiningOp();
41-
if (res_op) current_path.push_front(res_op);
42-
43-
constexpr int kNumExcludedOps = 2;
44-
collected_paths.push_back(RecurrenceCycle{
45-
operations: SmallVector<Operation *>(current_path.begin(), current_path.end()),
46-
length: static_cast<int>(current_path.size()) - kNumExcludedOps
47-
});
48-
49-
if (res_op) current_path.pop_front(); // Remove reserve before backtracking
50-
continue;
51-
}
52-
53-
if (Operation *def_op = operand.getDefiningOp()) {
54-
traverseAlongPath(def_op, reserve_value, current_path, visited_in_path, collected_paths);
55-
}
56-
}
57-
58-
current_path.pop_front(); // Backtrack
59-
visited_in_path.erase(op); // Unmark from path
60-
}
61-
62-
/// Collects all recurrence cycles rooted at reserve operations and closed by ctrl_mov.
63-
/// Each cycle contains the operation sequence and its corresponding length.
64-
SmallVector<RecurrenceCycle, 4> collectRecurrenceCycles(Operation *root_op) {
65-
SmallVector<RecurrenceCycle, 4> recurrence_cycles;
66-
67-
root_op->walk([&](neura::CtrlMovOp ctrl_mov_op) {
68-
Value target = ctrl_mov_op.getTarget();
69-
auto reserve_op = target.getDefiningOp<neura::ReserveOp>();
70-
if (!reserve_op)
71-
return;
72-
73-
Value reserve_value = reserve_op.getResult();
74-
Value ctrl_mov_from = ctrl_mov_op.getValue();
75-
76-
Operation *parent_op = ctrl_mov_from.getDefiningOp();
77-
if (!parent_op)
78-
return;
79-
80-
std::deque<Operation *> current_path;
81-
SmallVector<RecurrenceCycle, 4> collected_paths;
82-
DenseSet<Operation *> visited_in_path;
83-
traverseAlongPath(parent_op, reserve_value, current_path, visited_in_path, collected_paths);
84-
85-
for (auto &cycle : collected_paths) {
86-
cycle.operations.push_back(ctrl_mov_op);
87-
++cycle.length;
88-
recurrence_cycles.push_back(std::move(cycle));
89-
}
90-
});
91-
92-
return recurrence_cycles;
93-
}
94-
9521
struct MapToAcceleratorPass
9622
: public PassWrapper<MapToAcceleratorPass, OperationPass<ModuleOp>> {
9723
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MapToAcceleratorPass)
@@ -127,10 +53,17 @@ struct MapToAcceleratorPass
12753
<< longest->length << "):\n";
12854
for (Operation *op : longest->operations)
12955
op->print(llvm::errs()), llvm::errs() << "\n";
130-
IntegerAttr mii_attr = IntegerAttr::get(
56+
IntegerAttr rec_mii_attr = IntegerAttr::get(
13157
IntegerType::get(func.getContext(), 32), longest->length);
132-
func->setAttr("RecMII", mii_attr);
58+
func->setAttr("RecMII", rec_mii_attr);
13359
}
60+
61+
AcceleratorConfig config{/*numTiles=*/8}; // Example
62+
int res_mii = calculateResMii(func, config);
63+
IntegerAttr res_mii_attr = IntegerAttr::get(
64+
IntegerType::get(func.getContext(), 32), res_mii);
65+
func->setAttr("ResMII", res_mii_attr);
66+
13467
});
13568
}
13669
};

‎lib/NeuraDialect/Transforms/mapping/mapping_util.cpp‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
#include <deque>
2+
3+
#include "NeuraDialect/mapping/mapping_util.h"
4+
#include "NeuraDialect/NeuraOps.h"
5+
#include "mlir/Dialect/Func/IR/FuncOps.h"
6+
#include "mlir/IR/Operation.h"
7+
8+
using namespace mlir;
9+
using namespace mlir::neura;
10+
11+
namespace {
12+
13+
// Traverses (backward) the operation graph starting from the given operation
14+
// towards reserve_value.
15+
void traverseAlongPath(Operation *op, Value reserve_value,
16+
std::deque<Operation *> &current_path,
17+
DenseSet<Operation *> &visited_in_path,
18+
SmallVector<RecurrenceCycle, 4> &collected_paths) {
19+
if (!op || visited_in_path.contains(op))
20+
return;
21+
22+
visited_in_path.insert(op);
23+
current_path.push_front(op);
24+
25+
for (Value operand : op->getOperands()) {
26+
if (operand == reserve_value) {
27+
Operation *res_op = reserve_value.getDefiningOp();
28+
if (res_op) current_path.push_front(res_op);
29+
30+
constexpr int kNumExcludedOps = 2;
31+
collected_paths.push_back(RecurrenceCycle{
32+
operations: SmallVector<Operation *>(current_path.begin(), current_path.end()),
33+
length: static_cast<int>(current_path.size()) - kNumExcludedOps
34+
});
35+
36+
if (res_op) current_path.pop_front();
37+
continue;
38+
}
39+
40+
if (Operation *def_op = operand.getDefiningOp()) {
41+
traverseAlongPath(def_op, reserve_value, current_path, visited_in_path, collected_paths);
42+
}
43+
}
44+
45+
current_path.pop_front();
46+
visited_in_path.erase(op);
47+
}
48+
49+
} // namespace
50+
51+
SmallVector<RecurrenceCycle, 4> mlir::neura::collectRecurrenceCycles(Operation *func_op) {
52+
SmallVector<RecurrenceCycle, 4> recurrence_cycles;
53+
54+
func_op->walk([&](neura::CtrlMovOp ctrl_mov_op) {
55+
Value target = ctrl_mov_op.getTarget();
56+
auto reserve_op = target.getDefiningOp<neura::ReserveOp>();
57+
if (!reserve_op)
58+
return;
59+
60+
Value reserve_value = reserve_op.getResult();
61+
Value ctrl_mov_from = ctrl_mov_op.getValue();
62+
63+
Operation *parent_op = ctrl_mov_from.getDefiningOp();
64+
if (!parent_op)
65+
return;
66+
67+
std::deque<Operation *> current_path;
68+
SmallVector<RecurrenceCycle, 4> collected_paths;
69+
DenseSet<Operation *> visited_in_path;
70+
traverseAlongPath(parent_op, reserve_value, current_path, visited_in_path, collected_paths);
71+
72+
for (auto &cycle : collected_paths) {
73+
cycle.operations.push_back(ctrl_mov_op);
74+
++cycle.length;
75+
recurrence_cycles.push_back(std::move(cycle));
76+
}
77+
});
78+
79+
return recurrence_cycles;
80+
}
81+
82+
int mlir::neura::calculateResMii(Operation *func_op, const AcceleratorConfig &config) {
83+
int num_ops = 0;
84+
85+
// Count all "compute" operations (non-terminators, non-block ops).
86+
func_op->walk([&](Operation *op) {
87+
// Skips non-materialized ops.
88+
if (isa<func::FuncOp>(op) ||
89+
isa<neura::ConstantOp,
90+
neura::CtrlMovOp,
91+
neura::ReserveOp,
92+
neura::ReturnOp>(op)) {
93+
return;
94+
}
95+
++num_ops;
96+
});
97+
98+
llvm::errs() << "[calculateResMii] Total operations: " << num_ops << "\n";
99+
100+
// Avoid divide-by-zero
101+
int tiles = std::max(1, config.num_tiles);
102+
103+
return llvm::divideCeil(num_ops, tiles);
104+
}

‎test/neura/ctrl/branch_for.mlir‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
// RUN: --leverage-predicated-value \
1818
// RUN: --transform-ctrl-to-data-flow \
1919
// RUN: --map-to-accelerator \
20-
// RUN: | FileCheck %s -check-prefix=RECMII
20+
// RUN: | FileCheck %s -check-prefix=MII
2121

2222
func.func @loop_test() -> f32 {
2323
%n = llvm.mlir.constant(10 : i64) : i64
@@ -81,4 +81,4 @@ func.func @loop_test() -> f32 {
8181
// CTRL2DATA-NEXT: "neura.return"(%18) : (!neura.data<f32, i1>) -> ()
8282
// CTRL2DATA-NEXT: }
8383

84-
// RECMII: func.func @loop_test() -> f32 attributes {RecMII = 4 : i32, accelerator = "neura"}
84+
// MII: func.func @loop_test() -> f32 attributes {RecMII = 4 : i32, ResMII = 2 : i32, accelerator = "neura"}

0 commit comments

Comments
 (0)