1 |
(* translate-basis.sml |
(* translate-basis.sml |
2 |
* |
* |
3 |
* COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu) |
* COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu) |
4 |
* All rights reserved. |
* All rights reserved. |
5 |
* |
* |
6 |
* Translation for basis operations in Simple AST to HighIL code |
* Translation for basis operations in Simple AST to HighIL code |
12 |
* to the instantiated meta variables mvs) to a list of SSA assignments in |
* to the instantiated meta variables mvs) to a list of SSA assignments in |
13 |
* reverse order. |
* reverse order. |
14 |
*) |
*) |
15 |
val translate : (HighIL.var * Var.var * Types.meta_var list * HighIL.var list) |
val translate : (HighIL.var * Var.var * SimpleTypes.meta_arg list * HighIL.var list) |
16 |
-> HighIL.assign list |
-> HighIL.assignment list |
17 |
|
|
18 |
end = struct |
end = struct |
19 |
|
|
20 |
structure BV = BasisVars |
structure BV = BasisVars |
21 |
structure IL = HighIL |
structure IL = HighIL |
22 |
|
structure DstTy = HighILTypes |
23 |
structure Op = HighOps |
structure Op = HighOps |
24 |
structure Ty = Types |
structure Ty = SimpleTypes |
|
structure TU = TypeUtil |
|
|
structure MV = MetaVar |
|
25 |
structure VTbl = Var.Tbl |
structure VTbl = Var.Tbl |
26 |
|
|
27 |
fun assign (y, rator, xs) = [(y, IL.OP(rator, xs))] |
structure mk= mkOperators |
28 |
|
|
29 |
|
fun trType (Ty.TY ty) = TranslateTy.tr ty |
30 |
|
| trType _ = raise Fail "expected type" |
31 |
|
fun dimVarToInt (Ty.DIM d) = d |
32 |
|
| dimVarToInt _ = raise Fail "expected dim" |
33 |
|
fun dimVarToTensor dv = DstTy.tensorTy[dimVarToInt dv] |
34 |
|
fun dimVarToMatrix dv = let |
35 |
|
val d = dimVarToInt dv |
36 |
|
in |
37 |
|
DstTy.tensorTy[d, d] (* square matrix type *) |
38 |
|
end |
39 |
|
fun shapeVarToTensor (Ty.SHAPE shp) = DstTy.tensorTy shp |
40 |
|
| shapeVarToTensor _ = raise Fail "expected shape" |
41 |
|
|
42 |
|
fun assign (y, rator, xs) = [IL.ASSGN(y, IL.OP(rator, xs))] |
43 |
|
|
44 |
|
fun basisFn name (y, [], xs) = [IL.ASSGN(y, IL.APPLY(name, xs))] |
45 |
|
|
46 |
fun simpleOp rator (y, [], xs) = assign (y, rator, xs) |
fun simpleOp rator (y, [], xs) = assign (y, rator, xs) |
47 |
|
|
48 |
fun pruneDim d = (case TU.pruneDim d |
fun tensorOp rator (y, [sv], xs) = assign (y, rator(shapeVarToTensor sv), xs) |
49 |
of (Ty.DimConst n) => n |
|
50 |
| d => raise Fail("unresolved dimension " ^ TU.dimToString d) |
fun vectorOp rator (y, [dv], xs) = assign (y, rator(dimVarToTensor dv), xs) |
51 |
(* end case *)) |
|
52 |
|
fun kernel h (y, [], []) = assign(y, Op.Kernel(h, 0), []) |
53 |
|
|
54 |
|
(* utility functions for synthesizing eigenvector/eigenvalue code *) |
55 |
|
fun eigenVec (rator, dim) = let |
56 |
|
val ty = DstTy.SeqTy(DstTy.realTy, dim) |
57 |
|
in |
58 |
|
fn (y, _, [m]) => let |
59 |
|
val v = IL.Var.new("evals", ty) |
60 |
|
in |
61 |
|
[IL.MASSGN([v, y], rator, [m])] |
62 |
|
end |
63 |
|
end |
64 |
|
fun eigenVal (rator, dim) = let |
65 |
|
val ty = DstTy.SeqTy(DstTy.vecTy dim, dim) |
66 |
|
in |
67 |
|
fn (y, _, [m]) => let |
68 |
|
val v = IL.Var.new("evecs", ty) |
69 |
|
in |
70 |
|
[IL.MASSGN([y, v], rator, [m])] |
71 |
|
end |
72 |
|
end |
73 |
|
|
74 |
|
fun assignEin (y, rator, xs) = IL.ASSGN(y, IL.EINAPP(rator, xs)) |
75 |
|
|
76 |
|
fun simpleEOp rator (y, _,xs) = [assignEin(y, rator, xs)] |
77 |
|
|
|
fun tensorOp rator (y, [sv], xs) = (case TU.pruneShape(MV.toShape sv) |
|
|
of Ty.Shape dd => |
|
|
assign (y, rator(Op.TensorTy(List.map pruneDim dd)), xs) |
|
|
| shp => raise Fail("unresolved shape " ^ TU.shapeToString shp) |
|
|
(* end case *)) |
|
78 |
|
|
|
fun vectorOp rator (y, [dv], xs) = |
|
|
assign (y, rator(Op.TensorTy[pruneDim(MV.toDim dv)]), xs) |
|
79 |
|
|
80 |
fun kernel h (y, [], []) = assign(y, Op.Kernel h, []) |
|
81 |
|
|
82 |
|
(* shape is an int list, DIM is in int|variable, k-level of differntiation *) |
83 |
|
|
84 |
(* build a table that maps Basis variables to their translation functions *) |
(* build a table that maps Basis variables to their translation functions *) |
85 |
val tbl : ((IL.var * Ty.meta_var list * IL.var list) -> IL.assign list) VTbl.hash_table = let |
val tbl : ((IL.var * Ty.meta_arg list * IL.var list) -> IL.assignment list) VTbl.hash_table = let |
86 |
val tbl = VTbl.mkTable (128, Fail "Translate table") |
val tbl = VTbl.mkTable (128, Fail "Translate table") |
87 |
|
val insert = VTbl.insert tbl |
88 |
|
in |
89 |
|
List.app insert [ |
90 |
|
(BV.lt_ii, simpleOp(Op.LT DstTy.IntTy)), |
91 |
|
(BV.lt_rr, simpleOp(Op.LT DstTy.realTy)), |
92 |
|
(BV.lte_ii, simpleOp(Op.LTE DstTy.IntTy)), |
93 |
|
(BV.lte_rr, simpleOp(Op.LTE DstTy.realTy)), |
94 |
|
(BV.gte_ii, simpleOp(Op.GTE DstTy.IntTy)), |
95 |
|
(BV.gte_rr, simpleOp(Op.GTE(DstTy.realTy))), |
96 |
|
(BV.gt_ii, simpleOp(Op.GT DstTy.IntTy)), |
97 |
|
(BV.gt_rr, simpleOp(Op.GT(DstTy.realTy))), |
98 |
|
(BV.equ_bb, simpleOp(Op.EQ DstTy.BoolTy)), |
99 |
|
(BV.equ_ii, simpleOp(Op.EQ DstTy.IntTy)), |
100 |
|
(BV.equ_ss, simpleOp(Op.EQ DstTy.StringTy)), |
101 |
|
(BV.equ_rr, simpleOp(Op.EQ(DstTy.realTy))), |
102 |
|
(BV.neq_bb, simpleOp(Op.NEQ DstTy.BoolTy)), |
103 |
|
(BV.neq_ii, simpleOp(Op.NEQ DstTy.IntTy)), |
104 |
|
(BV.neq_ss, simpleOp(Op.NEQ DstTy.StringTy)), |
105 |
|
(BV.neq_rr, simpleOp(Op.NEQ(DstTy.realTy))), |
106 |
|
(BV.add_ii, simpleOp Op.IAdd), |
107 |
|
(BV.add_tt, fn (y, [shp], xs) => let |
108 |
|
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor shp |
109 |
|
val rator =mk.addTen(dd1) |
110 |
|
in |
111 |
|
[assignEin(y, rator,xs)] |
112 |
|
end), |
113 |
|
(BV.add_ff, fn (y, [_,Ty.DIM d, Ty.SHAPE dd], xs) => |
114 |
|
[assignEin(y, mk.addField(d, dd),xs)]), |
115 |
|
|
116 |
|
(BV.add_fr, fn (y, [_,Ty.DIM d], [f,s]) => |
117 |
|
[assignEin(y, mk.addTenField(d),[s,f])]), |
118 |
|
|
119 |
|
(BV.add_rf, fn (y, [_,Ty.DIM d], xs) => |
120 |
|
[assignEin(y, mk.addTenField(d),xs)]), |
121 |
|
|
122 |
|
|
123 |
|
(BV.sub_ii, simpleOp Op.ISub), |
124 |
|
(BV.sub_tt, fn (y, [shp], xs) => let |
125 |
|
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor shp |
126 |
|
val rator =mk.subTen(dd1) |
127 |
in |
in |
128 |
List.app (VTbl.insert tbl) [ |
[assignEin(y, rator,xs)] |
129 |
(BV.add_ii, simpleOp(Op.Add Op.IntTy)), |
end), |
130 |
(BV.add_tt, tensorOp Op.Add), |
(BV.sub_ff, fn (y, [_,Ty.DIM d, Ty.SHAPE dd], xs) => |
131 |
(BV.sub_ii, simpleOp(Op.Sub Op.IntTy)), |
[assignEin(y, mk.subField(d, dd),xs)]), |
132 |
(BV.sub_tt, tensorOp Op.Sub), |
|
133 |
(BV.mul_ii, simpleOp(Op.Mul Op.IntTy)), |
|
134 |
(BV.mul_rr, simpleOp(Op.Mul(Op.TensorTy[]))), |
(BV.sub_fr, fn (y, [_,Ty.DIM d], xs) => |
135 |
(BV.mul_rt, tensorOp Op.Scale), |
[assignEin(y, mk.subFieldTen(d),xs)]), |
136 |
(BV.mul_tr, fn (y, sv, [t, r]) => tensorOp Op.Scale (y, sv, [r, t])), |
|
137 |
(BV.div_ii, simpleOp(Op.Div Op.IntTy)), |
(BV.sub_rf, fn (y, [_,Ty.DIM d], xs) => |
138 |
(BV.div_rr, simpleOp(Op.Div(Op.TensorTy[]))), |
[assignEin(y, mk.subTenField(d),xs)]), |
139 |
(BV.div_tr, tensorOp Op.InvScale), |
|
140 |
(BV.lt_ii, simpleOp(Op.LT Op.IntTy)), |
(BV.mul_ii, simpleOp Op.IMul), |
141 |
(BV.lt_rr, simpleOp(Op.LT(Op.TensorTy[]))), |
(BV.mul_rr, fn (y,_,args) => [assignEin(y, mk.prodScalar,args)]), |
142 |
(BV.lte_ii, simpleOp(Op.LTE Op.IntTy)), |
(BV.mul_rt, fn (y, [shp], xs) => let |
143 |
(BV.lte_rr, simpleOp(Op.LTE(Op.TensorTy[]))), |
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor shp |
144 |
(BV.gte_ii, simpleOp(Op.GTE Op.IntTy)), |
val rator =mk.scaleTen(dd1) |
145 |
(BV.gte_rr, simpleOp(Op.GTE(Op.TensorTy[]))), |
in |
146 |
(BV.gt_ii, simpleOp(Op.GT Op.IntTy)), |
[assignEin(y, rator,xs)] |
147 |
(BV.gt_rr, simpleOp(Op.GT(Op.TensorTy[]))), |
end), |
148 |
(BV.equ_bb, simpleOp(Op.EQ Op.BoolTy)), |
(BV.mul_tr, fn (y, [shp], [t, r]) => let |
149 |
(BV.equ_ii, simpleOp(Op.EQ Op.IntTy)), |
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor shp |
150 |
(BV.equ_ss, simpleOp(Op.EQ Op.StringTy)), |
val rator = mk.scaleTen dd1 |
151 |
(BV.equ_rr, simpleOp(Op.EQ(Op.TensorTy[]))), |
in |
152 |
(BV.neq_bb, simpleOp(Op.NEQ Op.BoolTy)), |
[assignEin(y, rator,[r,t])] |
153 |
(BV.neq_ii, simpleOp(Op.NEQ Op.IntTy)), |
end), |
154 |
(BV.neq_ss, simpleOp(Op.NEQ Op.StringTy)), |
(BV.mul_rf, fn (y, [_,Ty.DIM d, Ty.SHAPE dd], xs) => |
155 |
(BV.neq_rr, simpleOp(Op.NEQ(Op.TensorTy[]))), |
[assignEin(y, mk.scaleField(d, dd),xs)]), |
156 |
(BV.neg_i, simpleOp(Op.Neg Op.IntTy)), |
(BV.mul_fr, fn (y, [_,Ty.DIM d, Ty.SHAPE dd], [f, s]) => |
157 |
(BV.neg_t, tensorOp Op.Neg), |
[assignEin(y, mk.scaleField(d, dd),[s,f])]), |
158 |
(BV.neg_f, fn (y, _, xs) => assign(y, Op.NegField, xs)), |
(BV.div_ii, simpleOp Op.IDiv), |
159 |
(BV.op_at, fn (y, _, xs) => assign(y, Op.Probe, xs)), |
(BV.div_rr, fn (y,_,args) => [assignEin(y, mk.divScalar,args)]), |
160 |
(BV.op_convolve, fn (y, _, xs) => assign(y, Op.Convolve, xs)), |
(BV.div_tr, fn (y, [shp], xs) => let |
161 |
(BV.op_D, fn (y, _, xs) => assign(y, Op.DiffField, xs)), |
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor shp |
162 |
(BV.op_norm, tensorOp Op.Norm), |
val rator =mk.divTen(dd1) |
163 |
|
in |
164 |
|
[assignEin(y, rator,xs)] |
165 |
|
end), |
166 |
|
(BV.div_fr, fn (y, [_,Ty.DIM d, Ty.SHAPE dd], xs) => |
167 |
|
[assignEin(y, mk.divideField(d, dd),xs)]), |
168 |
|
(BV.exp_ri, simpleOp(Op.Power)), |
169 |
|
(BV.exp_rr, basisFn MathFuns.pow), |
170 |
|
(BV.curl2D, simpleEOp mk.curl2d), |
171 |
|
(BV.curl3D, simpleEOp mk.curl3d), |
172 |
|
(BV.convolve_vk, fn (y, [_, Ty.DIM d, Ty.SHAPE dd], xs) => |
173 |
|
[assignEin(y, mk.conv(d, dd),xs)]), |
174 |
|
(BV.convolve_kv, fn (y, [_, Ty.DIM d, Ty.SHAPE dd], [k, v]) => |
175 |
|
[assignEin(y, mk.conv(d, dd),[v,k])]), |
176 |
|
(BV.neg_i, simpleOp Op.INeg), |
177 |
|
(BV.neg_t, fn (y, [shp], xs) => let |
178 |
|
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor shp |
179 |
|
val rator =mk.negTen(dd1) |
180 |
|
in |
181 |
|
[assignEin(y, rator,xs)] |
182 |
|
end), |
183 |
|
(BV.neg_f, fn (y, [_,Ty.DIM d, Ty.SHAPE dd], xs) => |
184 |
|
[assignEin(y, mk.negField(d, dd),xs)]), |
185 |
|
(BV.op_probe, fn (y, [_, Ty.DIM d, Ty.SHAPE dd], xs) => |
186 |
|
[assignEin(y, (mk.probe(dd,d,[d])),xs)]), |
187 |
|
(BV.op_D, fn (y, [_, Ty.DIM d], xs) => [assignEin(y, mk.grad([d]),xs)]), |
188 |
|
(BV.op_Dotimes, fn (y, [_, Ty.DIM d1, Ty.SHAPE dd, Ty.DIM d2], xs) => |
189 |
|
[assignEin(y, mk.divergence(d1, dd),xs)]), |
190 |
|
(BV.op_norm, fn (y, [sv], xs) => (case shapeVarToTensor sv |
191 |
|
of DstTy.TensorTy[] => assign(y, Op.Abs DstTy.realTy, xs) |
192 |
|
| ty => assign(y, Op.Norm ty, xs) |
193 |
|
(* end case *))), |
194 |
(BV.op_not, simpleOp Op.Not), |
(BV.op_not, simpleOp Op.Not), |
195 |
(* |
(BV.op_cross, simpleEOp mk.crossProduct), |
196 |
(BV.op_subscript, fn (y, [SK, NK], xs) => ??), (*FIXME*) |
(BV.op_outer, fn (y, [Ty.DIM d1, Ty.DIM d2], xs) => |
197 |
*) |
[assignEin(y, (mk.outerProduct(d1, d2)), xs)]), |
198 |
(BV.fn_CL, fn (y, _, xs) => assign(y, Op.CL, xs)), |
(BV.op_inner, fn (y, [sh1, sh2, _], xs) => let |
199 |
(BV.fn_convolve, fn (y, _, [h, img]) => assign(y, Op.Convolve, [img, h])), |
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor sh1 |
200 |
(BV.fn_cos, simpleOp Op.Cos), |
val ty2 as DstTy.TensorTy dd2 = shapeVarToTensor sh2 |
201 |
(BV.fn_dot, vectorOp Op.Dot), |
in |
202 |
(BV.fn_inside, fn (y, _, xs) => assign(y, Op.Inside, xs)), |
[assignEin(y, (mk.innerProduct(dd1,dd2)),xs)] |
203 |
(* |
end), |
204 |
(BV.fn_load, fn (y, [NK, SK], xs) => ??), (*FIXME*) |
(BV.op_colon, fn (y, [sh1, sh2, _], xs) => let |
205 |
*) |
val ty1 as DstTy.TensorTy dd1 = shapeVarToTensor sh1 |
206 |
(BV.fn_max, simpleOp Op.Min), |
val ty2 as DstTy.TensorTy dd2 = shapeVarToTensor sh2 |
207 |
(BV.fn_min, simpleOp Op.Max), |
in |
208 |
(BV.fn_modulate, vectorOp Op.Mul), |
[assignEin(y, (mk.doubleDot(dd1,dd2)),xs)] |
209 |
(BV.fn_pow, simpleOp Op.Pow), |
end), |
210 |
|
(BV.fn_inside, fn (y, [_, Ty.DIM d, _], xs) => assign(y, Op.Inside d, xs)), |
211 |
|
(BV.clamp_rrr, simpleOp (Op.Clamp DstTy.realTy)), |
212 |
|
(BV.clamp_vvv, vectorOp Op.Clamp), |
213 |
|
(BV.lerp3, tensorOp Op.Lerp), |
214 |
|
(BV.lerp5, fn (y, [sv], [a, b, x0, x, x1]) => let |
215 |
|
val t1 = IL.Var.new("t1", DstTy.realTy) |
216 |
|
val t2 = IL.Var.new("t2", DstTy.realTy) |
217 |
|
val t3 = IL.Var.new("t3", DstTy.realTy) |
218 |
|
in [ |
219 |
|
assignEin(t1, mk.subScalar,[x,x0]), |
220 |
|
assignEin(t2, mk.subScalar,[x1,x0]), |
221 |
|
assignEin(t3, mk.divScalar,[t1,t2]), |
222 |
|
IL.ASSGN(y, IL.OP(Op.Lerp(shapeVarToTensor sv), [a, b, t3])) |
223 |
|
] end), |
224 |
|
(BV.evals2x2, eigenVal (Op.Eigen2x2, 2)), |
225 |
|
(BV.evals3x3, eigenVal (Op.Eigen3x3, 3)), |
226 |
|
(BV.evecs2x2, eigenVec (Op.Eigen2x2, 2)), |
227 |
|
(BV.evecs3x3, eigenVec (Op.Eigen3x3, 3)), |
228 |
|
(BV.fn_max, simpleOp Op.Max), |
229 |
|
(BV.fn_min, simpleOp Op.Min), |
230 |
|
|
231 |
|
(* modulate is vector * vector pointwise multiplication *) |
232 |
|
(BV.fn_modulate, fn (y,[Ty.DIM dd1], xs) => |
233 |
|
[assignEin(y, (mk.modulate dd1),xs)]), |
234 |
|
(BV.fn_normalize, vectorOp Op.Normalize), |
235 |
(BV.fn_principleEvec, vectorOp Op.PrincipleEvec), |
(BV.fn_principleEvec, vectorOp Op.PrincipleEvec), |
236 |
(BV.fn_sin, simpleOp Op.Sin), |
(BV.fn_trace, fn (y, [Ty.DIM d], xs) => |
237 |
|
[assignEin(y,(mk.trace d), xs)]), |
238 |
|
(BV.fn_transpose, fn (y, [Ty.DIM d1, Ty.DIM d2], xs) => |
239 |
|
[assignEin(y, (mk.transpose [d1,d2]), xs)]), |
240 |
(BV.kn_bspln3, kernel Kernel.bspln3), |
(BV.kn_bspln3, kernel Kernel.bspln3), |
241 |
(BV.kn_bspln5, kernel Kernel.bspln5), |
(BV.kn_bspln5, kernel Kernel.bspln5), |
242 |
(BV.kn_ctmr, kernel Kernel.ctmr), |
(BV.kn_ctmr, kernel Kernel.ctmr), |
243 |
|
(BV.kn_c2ctmr, kernel Kernel.ctmr), |
244 |
|
(BV.kn_c4hexic, kernel Kernel.c4hexic), |
245 |
(BV.kn_tent, kernel Kernel.tent), |
(BV.kn_tent, kernel Kernel.tent), |
246 |
(BV.i2r, simpleOp Op.IntToReal)(*, |
(BV.kn_c1tent, kernel Kernel.tent), |
247 |
(BV.input, fn (y, [TK], xs) => ??), (*FIXME*) |
(BV.i2r, simpleOp Op.IntToReal), |
248 |
(BV.optInput, fn (y, [TK], xs) => ??) (*FIXME*) |
(BV.identity, fn (y, [Ty.DIM d], xs) => |
249 |
|
[assignEin(y, (mk.identity [d,d]), xs)]), |
250 |
|
(* (BV.zero, fn (y, [sv], []) => |
251 |
|
assign(y, Op.Zero(shapeVarToTensor sv), [])),*) |
252 |
|
(BV.subscript, fn (y, [ty, Ty.DIM d], xs) => |
253 |
|
assign (y, Op.SeqSub(DstTy.SeqTy(trType ty, d)), xs)) |
254 |
|
(*, |
255 |
|
(BV.dynSubscript, fn (y, [tv], args) => |
256 |
|
assign(y, Op.SeqSub(DstTy.DynSeqTy(pruneTy tv)), args)) |
257 |
*) |
*) |
258 |
]; |
]; |
259 |
|
(* add C math functions *) |
260 |
|
List.app (fn (n, x) => insert(x, basisFn n)) BV.mathFns; |
261 |
tbl |
tbl |
262 |
end |
end |
263 |
|
|
264 |
fun translate (y, f, mvs, xs) = (case VTbl.find tbl f |
fun translate (y, f, mvs, xs) = (case VTbl.find tbl f |
265 |
of SOME transFn => transFn(y, mvs, xs) |
of SOME transFn => transFn(y, mvs, xs) |
266 |
| NONE => raise Fail("TranslateBasis.translate: unknown function " ^ Var.uniqueNameOf f) |
| NONE => raise Fail("TranslateBasis.translate: unknown basis function " ^ Var.uniqueNameOf f) |
267 |
(* end case *)) |
(* end case *)) |
268 |
|
handle ex => (print(concat["translate (", IL.Var.toString y, ", ", |
269 |
|
Var.uniqueNameOf f, ", ...)\n"]); raise ex) |
270 |
|
|
271 |
end |
end |