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 |
32 |
(* end case *)) |
(* end case *)) |
33 |
|
|
34 |
fun pruneShape sv = (case TU.pruneShape(MV.toShape sv) |
fun pruneShape sv = (case TU.pruneShape(MV.toShape sv) |
35 |
of Ty.Shape dd => DstTy.TensorTy(List.map pruneDim dd) |
of Ty.Shape dd => DstTy.tensorTy(List.map pruneDim dd) |
36 |
| shp => raise Fail("unresolved shape " ^ TU.shapeToString shp) |
| shp => raise Fail("unresolved shape " ^ TU.shapeToString shp) |
37 |
(* end case *)) |
(* end case *)) |
38 |
|
|
39 |
fun dimVarToTensor dv = DstTy.TensorTy[pruneDim(MV.toDim dv)] |
fun dimVarToTensor dv = DstTy.tensorTy[pruneDim(MV.toDim dv)] |
40 |
|
fun dimVarToMatrix dv = let |
41 |
|
val d = pruneDim(MV.toDim dv) |
42 |
|
in |
43 |
|
DstTy.tensorTy[d, d] (* square matrix type *) |
44 |
|
end |
45 |
fun shapeVarToTensor sv = pruneShape sv |
fun shapeVarToTensor sv = pruneShape sv |
46 |
|
|
47 |
fun assign (y, rator, xs) = [(y, IL.OP(rator, xs))] |
fun assign (y, rator, xs) = [(y, IL.OP(rator, xs))] |
48 |
|
|
49 |
|
fun basisFn name (y, [], xs) = [(y, IL.APPLY(name, xs))] |
50 |
|
|
51 |
fun simpleOp rator (y, [], xs) = assign (y, rator, xs) |
fun simpleOp rator (y, [], xs) = assign (y, rator, xs) |
52 |
|
|
53 |
fun tensorOp rator (y, [sv], xs) = assign (y, rator(shapeVarToTensor sv), xs) |
fun tensorOp rator (y, [sv], xs) = assign (y, rator(shapeVarToTensor sv), xs) |
54 |
|
|
55 |
fun vectorOp rator (y, [dv], xs) = assign (y, rator(dimVarToTensor dv), xs) |
fun vectorOp rator (y, [dv], xs) = assign (y, rator(dimVarToTensor dv), xs) |
56 |
|
|
57 |
fun kernel h (y, [], []) = assign(y, Op.Kernel h, []) |
fun kernel h (y, [], []) = assign(y, Op.Kernel(h, 0), []) |
58 |
|
|
59 |
(* build a table that maps Basis variables to their translation functions *) |
(* build a table that maps Basis variables to their translation functions *) |
60 |
val tbl : ((IL.var * Ty.meta_var list * IL.var list) -> IL.assign list) VTbl.hash_table = let |
val tbl : ((IL.var * Ty.meta_var list * IL.var list) -> IL.assign list) VTbl.hash_table = let |
61 |
val tbl = VTbl.mkTable (128, Fail "Translate table") |
val tbl = VTbl.mkTable (128, Fail "Translate table") |
62 |
in |
in |
63 |
List.app (VTbl.insert tbl) [ |
List.app (VTbl.insert tbl) [ |
|
(BV.add_ii, simpleOp(Op.Add DstTy.IntTy)), |
|
|
(BV.add_tt, tensorOp Op.Add), |
|
|
(BV.sub_ii, simpleOp(Op.Sub DstTy.IntTy)), |
|
|
(BV.sub_tt, tensorOp Op.Sub), |
|
|
(BV.mul_ii, simpleOp(Op.Mul DstTy.IntTy)), |
|
|
(BV.mul_rr, simpleOp(Op.Mul(DstTy.TensorTy[]))), |
|
|
(BV.mul_rt, tensorOp Op.Scale), |
|
|
(BV.mul_tr, fn (y, sv, [t, r]) => tensorOp Op.Scale (y, sv, [r, t])), |
|
|
(BV.div_ii, simpleOp(Op.Div DstTy.IntTy)), |
|
|
(BV.div_rr, simpleOp(Op.Div(DstTy.TensorTy[]))), |
|
|
(BV.div_tr, tensorOp Op.InvScale), |
|
64 |
(BV.lt_ii, simpleOp(Op.LT DstTy.IntTy)), |
(BV.lt_ii, simpleOp(Op.LT DstTy.IntTy)), |
65 |
(BV.lt_rr, simpleOp(Op.LT(DstTy.TensorTy[]))), |
(BV.lt_rr, simpleOp(Op.LT DstTy.realTy)), |
66 |
(BV.lte_ii, simpleOp(Op.LTE DstTy.IntTy)), |
(BV.lte_ii, simpleOp(Op.LTE DstTy.IntTy)), |
67 |
(BV.lte_rr, simpleOp(Op.LTE(DstTy.TensorTy[]))), |
(BV.lte_rr, simpleOp(Op.LTE DstTy.realTy)), |
68 |
(BV.gte_ii, simpleOp(Op.GTE DstTy.IntTy)), |
(BV.gte_ii, simpleOp(Op.GTE DstTy.IntTy)), |
69 |
(BV.gte_rr, simpleOp(Op.GTE(DstTy.TensorTy[]))), |
(BV.gte_rr, simpleOp(Op.GTE(DstTy.realTy))), |
70 |
(BV.gt_ii, simpleOp(Op.GT DstTy.IntTy)), |
(BV.gt_ii, simpleOp(Op.GT DstTy.IntTy)), |
71 |
(BV.gt_rr, simpleOp(Op.GT(DstTy.TensorTy[]))), |
(BV.gt_rr, simpleOp(Op.GT(DstTy.realTy))), |
72 |
(BV.equ_bb, simpleOp(Op.EQ DstTy.BoolTy)), |
(BV.equ_bb, simpleOp(Op.EQ DstTy.BoolTy)), |
73 |
(BV.equ_ii, simpleOp(Op.EQ DstTy.IntTy)), |
(BV.equ_ii, simpleOp(Op.EQ DstTy.IntTy)), |
74 |
(BV.equ_ss, simpleOp(Op.EQ DstTy.StringTy)), |
(BV.equ_ss, simpleOp(Op.EQ DstTy.StringTy)), |
75 |
(BV.equ_rr, simpleOp(Op.EQ(DstTy.TensorTy[]))), |
(BV.equ_rr, simpleOp(Op.EQ(DstTy.realTy))), |
76 |
(BV.neq_bb, simpleOp(Op.NEQ DstTy.BoolTy)), |
(BV.neq_bb, simpleOp(Op.NEQ DstTy.BoolTy)), |
77 |
(BV.neq_ii, simpleOp(Op.NEQ DstTy.IntTy)), |
(BV.neq_ii, simpleOp(Op.NEQ DstTy.IntTy)), |
78 |
(BV.neq_ss, simpleOp(Op.NEQ DstTy.StringTy)), |
(BV.neq_ss, simpleOp(Op.NEQ DstTy.StringTy)), |
79 |
(BV.neq_rr, simpleOp(Op.NEQ(DstTy.TensorTy[]))), |
(BV.neq_rr, simpleOp(Op.NEQ(DstTy.realTy))), |
80 |
|
(BV.add_ii, simpleOp(Op.Add DstTy.IntTy)), |
81 |
|
(BV.add_tt, tensorOp Op.Add), |
82 |
|
(BV.add_ff, fn (y, _, [f, g]) => assign(y, Op.AddField, [f, g])), |
83 |
|
(BV.sub_ii, simpleOp(Op.Sub DstTy.IntTy)), |
84 |
|
(BV.sub_tt, tensorOp Op.Sub), |
85 |
|
(BV.sub_ff, fn (y, _, [f, g]) => assign(y, Op.SubField, [f, g])), |
86 |
|
(BV.mul_ii, simpleOp(Op.Mul DstTy.IntTy)), |
87 |
|
(BV.mul_rr, simpleOp(Op.Mul(DstTy.realTy))), |
88 |
|
(BV.mul_rt, tensorOp Op.Scale), |
89 |
|
(BV.mul_tr, fn (y, sv, [t, r]) => tensorOp Op.Scale (y, sv, [r, t])), |
90 |
|
(BV.mul_rf, fn (y, _, [s, f]) => assign(y, Op.ScaleField, [s, f])), |
91 |
|
(BV.mul_fr, fn (y, _, [f, s]) => assign(y, Op.ScaleField, [s, f])), |
92 |
|
(BV.div_ii, simpleOp(Op.Div DstTy.IntTy)), |
93 |
|
(BV.div_rr, simpleOp(Op.Div DstTy.realTy)), |
94 |
|
(BV.div_tr, fn (y, [sv], [x, s]) => let |
95 |
|
val one = IL.Var.new("one", DstTy.realTy) |
96 |
|
val s' = IL.Var.new("sInv", DstTy.realTy) |
97 |
|
in [ |
98 |
|
(one, IL.LIT(Literal.Float(FloatLit.one))), |
99 |
|
(s', IL.OP(Op.Div DstTy.realTy, [one, s])), |
100 |
|
(y, IL.OP(Op.Scale(shapeVarToTensor sv), [s', x])) |
101 |
|
] end), |
102 |
|
(BV.div_fr, fn (y, _, [f, s]) => let |
103 |
|
val one = IL.Var.new("one", DstTy.realTy) |
104 |
|
val s' = IL.Var.new("sInv", DstTy.realTy) |
105 |
|
in [ |
106 |
|
(one, IL.LIT(Literal.Float(FloatLit.one))), |
107 |
|
(s', IL.OP(Op.Div DstTy.realTy, [one, s])), |
108 |
|
(y, IL.OP(Op.ScaleField, [s', f])) |
109 |
|
] end), |
110 |
|
(BV.exp_ri, simpleOp(Op.Power)), |
111 |
|
(BV.exp_rr, basisFn ILBasis.pow), |
112 |
|
(BV.convolve_vk, fn (y, [_, dv, _], xs) => |
113 |
|
assign(y, Op.Field(pruneDim(MV.toDim dv)), xs)), |
114 |
|
(BV.convolve_kv, fn (y, [_, dv, _], [k, v]) => |
115 |
|
assign(y, Op.Field(pruneDim(MV.toDim dv)), [v, k])), |
116 |
(BV.neg_i, simpleOp(Op.Neg DstTy.IntTy)), |
(BV.neg_i, simpleOp(Op.Neg DstTy.IntTy)), |
117 |
(BV.neg_t, tensorOp Op.Neg), |
(BV.neg_t, tensorOp Op.Neg), |
118 |
(BV.neg_f, fn (y, _, xs) => assign(y, Op.NegField, xs)), |
(BV.neg_f, fn (y, _, xs) => assign(y, Op.NegField, xs)), |
119 |
(BV.op_at, fn (y, [_, dv, sv], xs) => |
(BV.op_at, fn (y, [_, dv, sv], xs) => |
120 |
assign(y, Op.Probe(shapeVarToTensor sv, dimVarToTensor dv), xs)), |
assign(y, Op.Probe(dimVarToTensor dv, shapeVarToTensor sv), xs)), |
|
(BV.op_convolve, fn (y, _, xs) => assign(y, Op.Convolve, xs)), |
|
121 |
(BV.op_D, fn (y, _, xs) => assign(y, Op.DiffField, xs)), |
(BV.op_D, fn (y, _, xs) => assign(y, Op.DiffField, xs)), |
122 |
(BV.op_norm, tensorOp Op.Norm), |
(BV.op_Dotimes, fn (y, _, xs) => assign(y, Op.DiffField, xs)), |
123 |
|
(BV.op_norm, fn (y, [sv], xs) => (case shapeVarToTensor sv |
124 |
|
of DstTy.TensorTy[] => assign(y, Op.Abs DstTy.realTy, xs) |
125 |
|
| ty => assign(y, Op.Norm ty, xs) |
126 |
|
(* end case *))), |
127 |
(BV.op_not, simpleOp Op.Not), |
(BV.op_not, simpleOp Op.Not), |
128 |
|
(BV.fn_atan2, basisFn ILBasis.atan2), |
129 |
(BV.fn_CL, fn (y, _, xs) => assign(y, Op.CL, xs)), |
(BV.fn_CL, fn (y, _, xs) => assign(y, Op.CL, xs)), |
130 |
(BV.fn_convolve, fn (y, _, [h, img]) => assign(y, Op.Convolve, [img, h])), |
(BV.fn_convolve, fn (y, [_, dv, _], [h, img]) => |
131 |
(BV.fn_cos, simpleOp Op.Cos), |
assign(y, Op.Field(pruneDim(MV.toDim dv)), [img, h])), |
132 |
|
(BV.fn_cos, basisFn ILBasis.cos), |
133 |
|
(BV.op_cross, simpleOp Op.Cross), |
134 |
|
(BV.fn_cross, simpleOp Op.Cross), |
135 |
|
(BV.op_outer, fn (y, [dv1, dv2], xs) => let |
136 |
|
val d1 = pruneDim(MV.toDim dv1) |
137 |
|
val d2 = pruneDim(MV.toDim dv2) |
138 |
|
in |
139 |
|
assign (y, Op.Outer(DstTy.tensorTy[d1, d2]), xs) |
140 |
|
end), |
141 |
|
(BV.fn_outer, fn (y, [dv1, dv2], xs) => let |
142 |
|
val d1 = pruneDim(MV.toDim dv1) |
143 |
|
val d2 = pruneDim(MV.toDim dv2) |
144 |
|
in |
145 |
|
assign (y, Op.Outer(DstTy.tensorTy[d1, d2]), xs) |
146 |
|
end), |
147 |
|
(BV.op_inner, fn (y, [sh1, sh2, _], xs) => let |
148 |
|
val ty1 as DstTy.TensorTy dd1 = pruneShape sh1 |
149 |
|
val ty2 as DstTy.TensorTy dd2 = pruneShape sh2 |
150 |
|
val rator = (case (dd1, dd2) |
151 |
|
of ([d], [d']) => Op.Dot ty1 |
152 |
|
| ([d1], [d1', d2]) => Op.MulVecMat ty2 |
153 |
|
| ([d1, d2], [d2']) => Op.MulMatVec ty1 |
154 |
|
| ([d1, d2], [d2', d3]) => Op.MulMatMat(ty1, ty2) |
155 |
|
| _ => raise Fail "unsupported inner-product type" |
156 |
|
(* end case *)) |
157 |
|
in |
158 |
|
assign (y, rator, xs) |
159 |
|
end), |
160 |
(BV.fn_dot, vectorOp Op.Dot), |
(BV.fn_dot, vectorOp Op.Dot), |
161 |
|
(BV.fn_evals, fn _ => raise Fail "evals not implemented yet"), (* FIXME *) |
162 |
|
(BV.fn_evecs, fn _ => raise Fail "evecs not implemented yet"), (* FIXME *) |
163 |
(BV.fn_inside, fn (y, [_, dv, _], xs) => |
(BV.fn_inside, fn (y, [_, dv, _], xs) => |
164 |
assign(y, Op.Inside(dimVarToTensor dv), xs)), |
assign(y, Op.Inside(pruneDim(MV.toDim dv)), xs)), |
165 |
|
(BV.clamp_rrr, simpleOp (Op.Clamp DstTy.realTy)), |
166 |
|
(BV.clamp_vvv, vectorOp Op.Clamp), |
167 |
|
(BV.lerp3, tensorOp Op.Lerp), |
168 |
|
(BV.lerp5, fn (y, [sv], [a, b, x0, x, x1]) => let |
169 |
|
val t1 = IL.Var.new("t1", DstTy.realTy) |
170 |
|
val t2 = IL.Var.new("t2", DstTy.realTy) |
171 |
|
val t3 = IL.Var.new("t3", DstTy.realTy) |
172 |
|
in [ |
173 |
|
(t1, IL.OP(Op.Sub DstTy.realTy, [x, x0])), |
174 |
|
(t2, IL.OP(Op.Sub DstTy.realTy, [x1, x0])), |
175 |
|
(t3, IL.OP(Op.Div DstTy.realTy, [t1, t2])), |
176 |
|
(y, IL.OP(Op.Lerp(shapeVarToTensor sv), [a, b, t3])) |
177 |
|
] end), |
178 |
(BV.fn_max, simpleOp Op.Max), |
(BV.fn_max, simpleOp Op.Max), |
179 |
(BV.fn_min, simpleOp Op.Min), |
(BV.fn_min, simpleOp Op.Min), |
180 |
(BV.fn_modulate, vectorOp Op.Mul), |
(BV.fn_modulate, vectorOp Op.Mul), |
181 |
(BV.fn_pow, simpleOp Op.Pow), |
(BV.fn_normalize, vectorOp Op.Normalize), |
182 |
|
(BV.fn_outer, fn _ => raise Fail "outer not implemented yet"), (* FIXME *) |
183 |
|
(BV.fn_pow, basisFn ILBasis.pow), |
184 |
(BV.fn_principleEvec, vectorOp Op.PrincipleEvec), |
(BV.fn_principleEvec, vectorOp Op.PrincipleEvec), |
185 |
(BV.fn_sin, simpleOp Op.Sin), |
(BV.fn_sin, basisFn ILBasis.sin), |
186 |
|
(BV.fn_sqrt, basisFn ILBasis.sqrt), |
187 |
|
(BV.fn_tan, basisFn ILBasis.tan), |
188 |
|
(BV.fn_trace, fn (y, [dv], xs) => assign(y, Op.Trace(dimVarToMatrix dv), xs)), |
189 |
(BV.kn_bspln3, kernel Kernel.bspln3), |
(BV.kn_bspln3, kernel Kernel.bspln3), |
190 |
(BV.kn_bspln5, kernel Kernel.bspln5), |
(BV.kn_bspln5, kernel Kernel.bspln5), |
191 |
(BV.kn_ctmr, kernel Kernel.ctmr), |
(BV.kn_ctmr, kernel Kernel.ctmr), |
192 |
|
(BV.kn_c2ctmr, kernel Kernel.ctmr), |
193 |
(BV.kn_tent, kernel Kernel.tent), |
(BV.kn_tent, kernel Kernel.tent), |
194 |
(BV.i2r, simpleOp Op.IntToReal) |
(BV.kn_c1tent, kernel Kernel.tent), |
195 |
|
(BV.i2r, simpleOp Op.IntToReal), |
196 |
|
(BV.identity, fn (y, [dv], []) => |
197 |
|
assign(y, Op.Identity(pruneDim(MV.toDim dv)), [])), |
198 |
|
(BV.zero, fn (y, [sv], []) => |
199 |
|
assign(y, Op.Zero(shapeVarToTensor sv), [])) |
200 |
]; |
]; |
201 |
tbl |
tbl |
202 |
end |
end |
205 |
of SOME transFn => transFn(y, mvs, xs) |
of SOME transFn => transFn(y, mvs, xs) |
206 |
| NONE => raise Fail("TranslateBasis.translate: unknown basis function " ^ Var.uniqueNameOf f) |
| NONE => raise Fail("TranslateBasis.translate: unknown basis function " ^ Var.uniqueNameOf f) |
207 |
(* end case *)) |
(* end case *)) |
208 |
|
handle ex => (print(concat["translate (", IL.Var.toString y, ", ", |
209 |
|
Var.uniqueNameOf f, ", ...)\n"]); raise ex) |
210 |
|
|
211 |
end |
end |