-
Notifications
You must be signed in to change notification settings - Fork 175
/
ConstProp.scala
295 lines (268 loc) · 10.5 KB
/
ConstProp.scala
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
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
// See LICENSE for license details.
package firrtl.passes
import firrtl._
import firrtl.ir._
import firrtl.Utils._
import firrtl.Mappers._
import firrtl.PrimOps._
import annotation.tailrec
object ConstProp extends Pass {
private def pad(e: Expression, t: Type) = (bitWidth(e.tpe), bitWidth(t)) match {
case (we, wt) if we < wt => DoPrim(Pad, Seq(e), Seq(wt), t)
case (we, wt) if we == wt => e
}
private def asUInt(e: Expression, t: Type) = DoPrim(AsUInt, Seq(e), Seq(), t)
trait FoldLogicalOp {
def fold(c1: Literal, c2: Literal): Expression
def simplify(e: Expression, lhs: Literal, rhs: Expression): Expression
def apply(e: DoPrim): Expression = (e.args.head, e.args(1)) match {
case (lhs: Literal, rhs: Literal) => fold(lhs, rhs)
case (lhs: Literal, rhs) => pad(simplify(e, lhs, rhs), e.tpe)
case (lhs, rhs: Literal) => pad(simplify(e, rhs, lhs), e.tpe)
case _ => e
}
}
object FoldAND extends FoldLogicalOp {
def fold(c1: Literal, c2: Literal) = UIntLiteral(c1.value & c2.value, c1.width max c2.width)
def simplify(e: Expression, lhs: Literal, rhs: Expression) = lhs match {
case UIntLiteral(v, w) if v == BigInt(0) => UIntLiteral(0, w)
case SIntLiteral(v, w) if v == BigInt(0) => UIntLiteral(0, w)
case UIntLiteral(v, IntWidth(w)) if v == (BigInt(1) << bitWidth(rhs.tpe).toInt) - 1 => rhs
case _ => e
}
}
object FoldOR extends FoldLogicalOp {
def fold(c1: Literal, c2: Literal) = UIntLiteral(c1.value | c2.value, c1.width max c2.width)
def simplify(e: Expression, lhs: Literal, rhs: Expression) = lhs match {
case UIntLiteral(v, _) if v == BigInt(0) => rhs
case SIntLiteral(v, _) if v == BigInt(0) => asUInt(rhs, e.tpe)
case UIntLiteral(v, IntWidth(w)) if v == (BigInt(1) << bitWidth(rhs.tpe).toInt) - 1 => lhs
case _ => e
}
}
object FoldXOR extends FoldLogicalOp {
def fold(c1: Literal, c2: Literal) = UIntLiteral(c1.value ^ c2.value, c1.width max c2.width)
def simplify(e: Expression, lhs: Literal, rhs: Expression) = lhs match {
case UIntLiteral(v, _) if v == BigInt(0) => rhs
case SIntLiteral(v, _) if v == BigInt(0) => asUInt(rhs, e.tpe)
case _ => e
}
}
object FoldEqual extends FoldLogicalOp {
def fold(c1: Literal, c2: Literal) = UIntLiteral(if (c1.value == c2.value) 1 else 0, IntWidth(1))
def simplify(e: Expression, lhs: Literal, rhs: Expression) = lhs match {
case UIntLiteral(v, IntWidth(w)) if v == BigInt(1) && w == BigInt(1) && bitWidth(rhs.tpe) == BigInt(1) => rhs
case _ => e
}
}
object FoldNotEqual extends FoldLogicalOp {
def fold(c1: Literal, c2: Literal) = UIntLiteral(if (c1.value != c2.value) 1 else 0, IntWidth(1))
def simplify(e: Expression, lhs: Literal, rhs: Expression) = lhs match {
case UIntLiteral(v, IntWidth(w)) if v == BigInt(0) && w == BigInt(1) && bitWidth(rhs.tpe) == BigInt(1) => rhs
case _ => e
}
}
private def foldConcat(e: DoPrim) = (e.args.head, e.args(1)) match {
case (UIntLiteral(xv, IntWidth(xw)), UIntLiteral(yv, IntWidth(yw))) => UIntLiteral(xv << yw.toInt | yv, IntWidth(xw + yw))
case _ => e
}
private def foldShiftLeft(e: DoPrim) = e.consts.head.toInt match {
case 0 => e.args.head
case x => e.args.head match {
case UIntLiteral(v, IntWidth(w)) => UIntLiteral(v << x, IntWidth(w + x))
case SIntLiteral(v, IntWidth(w)) => SIntLiteral(v << x, IntWidth(w + x))
case _ => e
}
}
private def foldShiftRight(e: DoPrim) = e.consts.head.toInt match {
case 0 => e.args.head
case x => e.args.head match {
// TODO when amount >= x.width, return a zero-width wire
case UIntLiteral(v, IntWidth(w)) => UIntLiteral(v >> x, IntWidth((w - x) max 1))
// take sign bit if shift amount is larger than arg width
case SIntLiteral(v, IntWidth(w)) => SIntLiteral(v >> x, IntWidth((w - x) max 1))
case _ => e
}
}
private def foldComparison(e: DoPrim) = {
def foldIfZeroedArg(x: Expression): Expression = {
def isUInt(e: Expression): Boolean = e.tpe match {
case UIntType(_) => true
case _ => false
}
def isZero(e: Expression) = e match {
case UIntLiteral(value, _) => value == BigInt(0)
case SIntLiteral(value, _) => value == BigInt(0)
case _ => false
}
x match {
case DoPrim(Lt, Seq(a,b),_,_) if isUInt(a) && isZero(b) => zero
case DoPrim(Leq, Seq(a,b),_,_) if isZero(a) && isUInt(b) => one
case DoPrim(Gt, Seq(a,b),_,_) if isZero(a) && isUInt(b) => zero
case DoPrim(Geq, Seq(a,b),_,_) if isUInt(a) && isZero(b) => one
case ex => ex
}
}
def foldIfOutsideRange(x: Expression): Expression = {
//Note, only abides by a partial ordering
case class Range(min: BigInt, max: BigInt) {
def === (that: Range) =
Seq(this.min, this.max, that.min, that.max)
.sliding(2,1)
.map(x => x.head == x(1))
.reduce(_ && _)
def > (that: Range) = this.min > that.max
def >= (that: Range) = this.min >= that.max
def < (that: Range) = this.max < that.min
def <= (that: Range) = this.max <= that.min
}
def range(e: Expression): Range = e match {
case UIntLiteral(value, _) => Range(value, value)
case SIntLiteral(value, _) => Range(value, value)
case _ => e.tpe match {
case SIntType(IntWidth(width)) => Range(
min = BigInt(0) - BigInt(2).pow(width.toInt - 1),
max = BigInt(2).pow(width.toInt - 1) - BigInt(1)
)
case UIntType(IntWidth(width)) => Range(
min = BigInt(0),
max = BigInt(2).pow(width.toInt) - BigInt(1)
)
}
}
// Calculates an expression's range of values
x match {
case ex: DoPrim =>
def r0 = range(ex.args.head)
def r1 = range(ex.args(1))
ex.op match {
// Always true
case Lt if r0 < r1 => one
case Leq if r0 <= r1 => one
case Gt if r0 > r1 => one
case Geq if r0 >= r1 => one
// Always false
case Lt if r0 >= r1 => zero
case Leq if r0 > r1 => zero
case Gt if r0 <= r1 => zero
case Geq if r0 < r1 => zero
case _ => ex
}
case ex => ex
}
}
foldIfZeroedArg(foldIfOutsideRange(e))
}
private def constPropPrim(e: DoPrim): Expression = e.op match {
case Shl => foldShiftLeft(e)
case Shr => foldShiftRight(e)
case Cat => foldConcat(e)
case And => FoldAND(e)
case Or => FoldOR(e)
case Xor => FoldXOR(e)
case Eq => FoldEqual(e)
case Neq => FoldNotEqual(e)
case (Lt | Leq | Gt | Geq) => foldComparison(e)
case Not => e.args.head match {
case UIntLiteral(v, IntWidth(w)) => UIntLiteral(v ^ ((BigInt(1) << w.toInt) - 1), IntWidth(w))
case _ => e
}
case AsUInt => e.args.head match {
case SIntLiteral(v, IntWidth(w)) => UIntLiteral(v + (if (v < 0) BigInt(1) << w.toInt else 0), IntWidth(w))
case u: UIntLiteral => u
case _ => e
}
case AsSInt => e.args.head match {
case UIntLiteral(v, IntWidth(w)) => SIntLiteral(v - ((v >> (w.toInt-1)) << w.toInt), IntWidth(w))
case s: SIntLiteral => s
case _ => e
}
case Pad => e.args.head match {
case UIntLiteral(v, IntWidth(w)) => UIntLiteral(v, IntWidth(e.consts.head max w))
case SIntLiteral(v, IntWidth(w)) => SIntLiteral(v, IntWidth(e.consts.head max w))
case _ if bitWidth(e.args.head.tpe) == e.consts.head => e.args.head
case _ => e
}
case Bits => e.args.head match {
case lit: Literal =>
val hi = e.consts.head.toInt
val lo = e.consts(1).toInt
require(hi >= lo)
UIntLiteral((lit.value >> lo) & ((BigInt(1) << (hi - lo + 1)) - 1), getWidth(e.tpe))
case x if bitWidth(e.tpe) == bitWidth(x.tpe) => x.tpe match {
case t: UIntType => x
case _ => asUInt(x, e.tpe)
}
case _ => e
}
case _ => e
}
private def constPropMuxCond(m: Mux) = m.cond match {
case UIntLiteral(c, _) => pad(if (c == BigInt(1)) m.tval else m.fval, m.tpe)
case _ => m
}
private def constPropMux(m: Mux): Expression = (m.tval, m.fval) match {
case _ if m.tval == m.fval => m.tval
case (t: UIntLiteral, f: UIntLiteral) =>
if (t.value == BigInt(1) && f.value == BigInt(0) && bitWidth(m.tpe) == BigInt(1)) m.cond
else constPropMuxCond(m)
case _ => constPropMuxCond(m)
}
private def constPropNodeRef(r: WRef, e: Expression) = e match {
case _: UIntLiteral | _: SIntLiteral | _: WRef => e
case _ => r
}
// Two pass process
// 1. Propagate constants in expressions and forward propagate references
// 2. Propagate references again for backwards reference (Wires)
// TODO Replacing all wires with nodes makes the second pass unnecessary
@tailrec
private def constPropModule(m: Module): Module = {
var nPropagated = 0L
val nodeMap = collection.mutable.HashMap[String, Expression]()
def backPropExpr(expr: Expression): Expression = {
val old = expr map backPropExpr
val propagated = old match {
case ref @ WRef(rname, _,_, MALE) if nodeMap.contains(rname) =>
constPropNodeRef(ref, nodeMap(rname))
case x => x
}
if (old ne propagated) {
nPropagated += 1
}
propagated
}
def backPropStmt(stmt: Statement): Statement = stmt map backPropStmt map backPropExpr
def constPropExpression(e: Expression): Expression = {
val old = e map constPropExpression
val propagated = old match {
case p: DoPrim => constPropPrim(p)
case m: Mux => constPropMux(m)
case ref @ WRef(rname, _,_, MALE) if nodeMap.contains(rname) =>
constPropNodeRef(ref, nodeMap(rname))
case x => x
}
propagated
}
def constPropStmt(s: Statement): Statement = {
val stmtx = s map constPropStmt map constPropExpression
stmtx match {
case x: DefNode => nodeMap(x.name) = x.value
case Connect(_, WRef(wname, wtpe, WireKind, _), expr) =>
val exprx = constPropExpression(pad(expr, wtpe))
nodeMap(wname) = exprx
case _ =>
}
stmtx
}
val res = Module(m.info, m.name, m.ports, backPropStmt(constPropStmt(m.body)))
if (nPropagated > 0) constPropModule(res) else res
}
def run(c: Circuit): Circuit = {
val modulesx = c.modules.map {
case m: ExtModule => m
case m: Module => constPropModule(m)
}
Circuit(c.info, modulesx, c.main)
}
}