1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
use std::fmt::Debug;
use zhc_utils::{
iter::{CollectInSmallVec, CollectInVec},
small::SmallVec,
svec,
};
use crate::{
AnnIR, IR, OpMap, ValMap,
visualization::{LayoutDialect, LayoutInstructionSet},
};
use super::*;
#[derive(Clone, PartialEq, Eq)]
pub enum PlacementVariable {
NonGroup {
// The place cell for the op in question.
// That is it contains a value between zero and the length of the layer.
op: PlaceCell,
// The place cells for the args of the op.
// That is it contains values between zero and the total number of args of the ops of the
// layer. Follows the order of the args themselves.
args: SmallVec<PlaceCell>,
// The place cells for the rets of the op.
// That is it contains values between zero and the total number of returns of the ops of
// the layer. Follows the order of the rets themselves.
rets: SmallVec<PlaceCell>,
},
Group {
// The place cell for the op in question.
// That is it contains a value between zero and the length of the layer.
op: PlaceCell,
// The place cells for the args of the op.
// That is it contains values between zero and the total number of args of the ops of the
// layer. Follows the order of the args themselves.
args: SmallVec<PlaceCell>,
// The place cells for the rets of the op.
// That is it contains values between zero and the total number of returns of the ops of
// the layer. Follows the order of the rets themselves.
rets: SmallVec<PlaceCell>,
// The place cells for the input ops of the group.
// Note that those are also stored as op of NonGroup variant somewhere else.
// This is an alias that is used to simplify the bridge with the content of the group when
// recursing. The values are between zero and number of inputs of the group.
// Follows the order of the args.
inputs: SmallVec<PlaceCell>,
// The place cells for the outputs ops of the group.
// Note that those are also stored as op of NonGroup variant somewhere else.
// This is an alias that is used to simplify the bridge with the content of the group when
// recursing. The values are between zero and number of inputs of the group.
// Follows the order of the args.
outputs: SmallVec<PlaceCell>,
// The maps to pass annotations down to the nested ir.
maps: (OpMap<PlacementVariable>, ValMap<()>),
},
}
impl Debug for PlacementVariable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NonGroup { op, args, rets } => {
write!(f, "op: {op}, args: {args}, rets: {rets}")
}
Self::Group {
op,
args,
rets,
inputs,
outputs,
..
} => {
write!(
f,
"op: {op}, args: {args}, rets: {rets}, inputs: {inputs}, outputs: {outputs}"
)
}
}
}
}
impl PlacementVariable {
pub fn get_op_position(&self) -> &PlaceCell {
match self {
PlacementVariable::NonGroup { op, .. } | PlacementVariable::Group { op, .. } => op,
}
}
pub fn get_arg_positions(&self) -> &[PlaceCell] {
match self {
PlacementVariable::NonGroup { args, .. } | PlacementVariable::Group { args, .. } => {
args.as_slice()
}
}
}
pub fn get_ret_positions(&self) -> &[PlaceCell] {
match self {
PlacementVariable::NonGroup { rets, .. } | PlacementVariable::Group { rets, .. } => {
rets.as_slice()
}
}
}
}
pub fn annotate_for_solving(
ir: &IR<LayoutDialect>,
) -> AnnIR<'_, LayoutDialect, PlacementVariable, ()> {
ir.forward_dataflow_analysis(|opref| {
use LayoutInstructionSet::*;
let opann = match opref.get_instruction() {
Operation { .. } | Dummy { .. } | GroupInput { .. } | GroupOutput { .. } => {
PlacementVariable::NonGroup {
op: PlaceCell::new(0),
args: std::iter::repeat_with(|| PlaceCell::new(0))
.take(opref.get_args_arity())
.collect(),
rets: std::iter::repeat_with(|| PlaceCell::new(0))
.take(opref.get_return_arity())
.collect(),
}
}
Group { ir, .. } => {
let (opmap, valmap) = annotate_for_solving(&ir).into_maps();
let mut inputs = ir
.walk_ops_linear()
.filter(|a| {
matches!(a.get_instruction(), LayoutInstructionSet::GroupInput { .. })
})
.covec();
inputs.sort_unstable_by_key(|op| {
let LayoutInstructionSet::GroupInput { pos, .. } = op.get_instruction() else {
unreachable!()
};
pos
});
let inputs = inputs
.into_iter()
.map(|op| opmap.get(&*op).unwrap().get_op_position().clone())
.cosvec();
let mut outputs = ir
.walk_ops_linear()
.filter(|a| {
matches!(
a.get_instruction(),
LayoutInstructionSet::GroupOutput { .. }
)
})
.covec();
outputs.sort_unstable_by_key(|op| {
let LayoutInstructionSet::GroupOutput { pos, .. } = op.get_instruction() else {
unreachable!()
};
pos
});
let outputs = outputs
.into_iter()
.map(|op| opmap.get(&*op).unwrap().get_op_position().clone())
.cosvec();
let op = PlaceCell::new(0);
let args = std::iter::repeat_with(|| PlaceCell::new(0))
.take(opref.get_args_arity())
.collect();
let rets = std::iter::repeat_with(|| PlaceCell::new(0))
.take(opref.get_return_arity())
.collect();
PlacementVariable::Group {
op,
args,
rets,
inputs,
outputs,
maps: (opmap, valmap),
}
}
};
(opann, svec![(); opref.get_return_arity()])
})
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PlacementSolution {
NonGroup {
op: Place,
},
Group {
op: Place,
inputs: SmallVec<Place>,
outputs: SmallVec<Place>,
maps: (OpMap<PlacementSolution>, ValMap<()>),
},
}
impl PlacementSolution {
pub fn get_place(&self) -> Place {
match self {
PlacementSolution::NonGroup { op } | PlacementSolution::Group { op, .. } => *op,
}
}
}
fn resolve(input: OpMap<PlacementVariable>) -> OpMap<PlacementSolution> {
input.map(|v| match v {
PlacementVariable::NonGroup { op, .. } => PlacementSolution::NonGroup { op: op.get_val() },
PlacementVariable::Group {
op,
maps,
inputs,
outputs,
..
} => PlacementSolution::Group {
op: op.get_val(),
inputs: inputs.into_iter().map(|a| a.get_val()).collect(),
outputs: outputs.into_iter().map(|a| a.get_val()).collect(),
maps: (resolve(maps.0), maps.1.map(|_| ())),
},
})
}
pub fn turn_to_solution(
ir: AnnIR<'_, LayoutDialect, PlacementVariable, ()>,
) -> AnnIR<'_, LayoutDialect, PlacementSolution, ()> {
let AnnIR {
ir,
op_annotations,
val_annotations,
} = ir;
let op_annotations = resolve(op_annotations);
AnnIR::new(ir, op_annotations, val_annotations.map(|_| ()))
}