forked from mit-plv/koika
-
Notifications
You must be signed in to change notification settings - Fork 1
/
fft.v
301 lines (264 loc) · 16.8 KB
/
fft.v
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
296
297
298
299
300
301
(*! Computing an FFT !*)
Require Import Koika.Frontend.
Section KArray.
Definition array_init T n {reg_t ext_fn_t} : UInternalFunction reg_t ext_fn_t :=
{{ fun array_init (X: T ) : array_t {| array_type:= T; array_len:= n|} =>
`USugar( UArrayInit T (repeat {{X}} n) )`
}}.
Definition static_access {reg_t ext_fn_t} array idx: uaction reg_t ext_fn_t :=
(UUnop (UArray1 (UGetElement idx)) array).
Definition static_update {reg_t ext_fn_t} array idx v: uaction reg_t ext_fn_t :=
(UBinop (UArray2 (USubstElement idx)) array v).
Definition static_slice {reg_t ext_fn_t} (T: type) (init: uaction reg_t ext_fn_t) sz array idx: uaction reg_t ext_fn_t :=
let ainit := @array_init T sz reg_t ext_fn_t in
{{
let returned := ainit(`init`) in
`List.fold_left (fun acc el =>
static_update acc el ( UUnop (UArray1 (UGetElement (idx+el))) array)) (seq 0 sz) {{returned}}`
}}.
Definition access n T {reg_t} : UInternalFunction reg_t empty_ext_fn_t :=
let sz := log2 n in
{{ fun access (array : array_t {| array_type:= T; array_len:= n|} ) (idx: bits_t sz ) : T =>
` (USugar (USwitch {{idx}}
(USugar (UConstBits (sz:= sz) Bits.zero))
(List.map (fun branch =>
((USugar (UConstBits (Bits.of_nat (sz<:nat) branch))),
(UUnop (UArray1 (UGetElement branch)) ) {{array}}))
(seq 0 n)))) ` }}.
Definition update n T {reg_t} : UInternalFunction reg_t empty_ext_fn_t :=
let sz := log2 n in
{{ fun update (array : array_t {| array_type:= T; array_len:= n|} ) (idx: bits_t sz ) (v: T) : array_t {| array_type:= T; array_len:= n|} =>
` (USugar (USwitch {{idx}}
(USugar (UConstBits (sz:= sz) Bits.zero))
(List.map (fun branch =>
((USugar (UConstBits (Bits.of_nat (sz<:nat) branch))),
(UBinop (UArray2 (USubstElement branch)) ) {{array}} {{v}}))
(seq 0 n)))) ` }}.
End KArray.
Section FFT.
Definition complex :=
{| struct_name := "complex";
struct_fields :=
[ ("re", bits_t 32); ("im", bits_t 32) ] |}.
Definition log_points := 4.
Definition points := pow2 log_points.
Inductive reg_t := localSt.
Notation ext_fn_t := empty_ext_fn_t.
(* Inductive ext_fn_t := *)
(* foo *)
(* . *)
(* Definition Sigma (fn: ext_fn_t) := *)
(* match fn with *)
(* | foo => {$ bits_t 1 ~> bits_t 1 $} *)
(* end. *)
Notation Sigma := empty_Sigma.
Definition ctimes : UInternalFunction reg_t ext_fn_t :=
{{ fun ctimes (a: struct_t complex) (b: struct_t complex) : struct_t complex =>
let re := (get(a, re) * get(b,re))[#(Bits.of_nat 6 16) :+ 32] - (get(b,im) * get(a,im))[#(Bits.of_nat 6 16) :+ 32] in
let im := (get(a,re) * get(b,im))[#(Bits.of_nat 6 16) :+ 32] + (get(a,im) * get(b,re))[#(Bits.of_nat 6 16) :+ 32] in
struct complex { re := re; im:= im }
}}.
Definition cplus : UInternalFunction reg_t ext_fn_t :=
{{ fun ctimes (a: struct_t complex) (b: struct_t complex) : struct_t complex =>
let re := get(a, re) + get(b,re) in
let im := get(b,im) + get(a,im)in
struct complex { re := re; im:= im }
}}.
Definition csub : UInternalFunction reg_t ext_fn_t :=
{{ fun ctimes (a: struct_t complex) (b: struct_t complex) : struct_t complex =>
let re := get(a, re) - get(b,re) in
let im := get(a,im) - get(b,im)in
struct complex { re := re; im:= im }
}}.
Definition bfly2 : UInternalFunction reg_t ext_fn_t :=
let ainit := array_init (reg_t:= reg_t) (struct_t complex) 2 in
{{ fun bfly2
(t: array_t {| array_type := struct_t complex ; array_len:= 2|})
(k: struct_t complex) :
array_t {| array_type := struct_t complex ; array_len:= 2|} =>
let m := ctimes(`static_access {{t}} 1`, k) in
let zero := #(Bits.of_nat 32 0) in
let zerocomplex := struct complex {re:= zero; im := zero} in
let result := ainit(zerocomplex) in
(set result := `static_update {{result}} 0 {{cplus(`static_access {{t}} 0`,m)}}`);
`static_update {{result}} 1 {{csub(`static_access {{t}} 0`, m)}}`
}}.
Definition permute (dst:nat) (points:nat) : nat :=
if Nat.ltb dst (points/2) then
dst*2
else
(dst - points/2)*2 + 1.
Definition twiddle (stage:nat) (idx:nat) : uaction reg_t ext_fn_t := match stage, idx with
| 0, 0 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 1 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 2 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 3 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 4 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 5 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 6 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 7 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 8 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 9 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 10 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 11 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 12 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 13 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 14 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 0, 15 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 0 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 1 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 2 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 3 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 4 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 4294901760) in struct complex {re:=r;im:=i} }}
| 1, 5 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 4294901760) in struct complex {re:=r;im:=i} }}
| 1, 6 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 4294901760) in struct complex {re:=r;im:=i} }}
| 1, 7 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 4294901760) in struct complex {re:=r;im:=i} }}
| 1, 8 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 9 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 10 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 11 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 1, 12 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 65536) in struct complex {re:=r;im:=i} }}
| 1, 13 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 65536) in struct complex {re:=r;im:=i} }}
| 1, 14 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 65536) in struct complex {re:=r;im:=i} }}
| 1, 15 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 65536) in struct complex {re:=r;im:=i} }}
| 2, 0 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 2, 1 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 2, 2 => {{let r := #(Bits.of_N 32 46340) in let i := #(Bits.of_N 32 4294920955) in struct complex {re:=r;im:=i} }}
| 2, 3 => {{let r := #(Bits.of_N 32 46340) in let i := #(Bits.of_N 32 4294920955) in struct complex {re:=r;im:=i} }}
| 2, 4 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 4294901760) in struct complex {re:=r;im:=i} }}
| 2, 5 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 4294901760) in struct complex {re:=r;im:=i} }}
| 2, 6 => {{let r := #(Bits.of_N 32 4294920955) in let i := #(Bits.of_N 32 4294920955) in struct complex {re:=r;im:=i} }}
| 2, 7 => {{let r := #(Bits.of_N 32 4294920955) in let i := #(Bits.of_N 32 4294920955) in struct complex {re:=r;im:=i} }}
| 2, 8 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 2, 9 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 2, 10 => {{let r := #(Bits.of_N 32 4294920955) in let i := #(Bits.of_N 32 46340) in struct complex {re:=r;im:=i} }}
| 2, 11 => {{let r := #(Bits.of_N 32 4294920955) in let i := #(Bits.of_N 32 46340) in struct complex {re:=r;im:=i} }}
| 2, 12 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 65536) in struct complex {re:=r;im:=i} }}
| 2, 13 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 65536) in struct complex {re:=r;im:=i} }}
| 2, 14 => {{let r := #(Bits.of_N 32 46340) in let i := #(Bits.of_N 32 46340) in struct complex {re:=r;im:=i} }}
| 2, 15 => {{let r := #(Bits.of_N 32 46340) in let i := #(Bits.of_N 32 46340) in struct complex {re:=r;im:=i} }}
| 3, 0 => {{let r := #(Bits.of_N 32 65536) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 3, 1 => {{let r := #(Bits.of_N 32 60547) in let i := #(Bits.of_N 32 4294942216) in struct complex {re:=r;im:=i} }}
| 3, 2 => {{let r := #(Bits.of_N 32 46340) in let i := #(Bits.of_N 32 4294920955) in struct complex {re:=r;im:=i} }}
| 3, 3 => {{let r := #(Bits.of_N 32 25079) in let i := #(Bits.of_N 32 4294906749) in struct complex {re:=r;im:=i} }}
| 3, 4 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 4294901760) in struct complex {re:=r;im:=i} }}
| 3, 5 => {{let r := #(Bits.of_N 32 4294942216) in let i := #(Bits.of_N 32 4294906749) in struct complex {re:=r;im:=i} }}
| 3, 6 => {{let r := #(Bits.of_N 32 4294920955) in let i := #(Bits.of_N 32 4294920955) in struct complex {re:=r;im:=i} }}
| 3, 7 => {{let r := #(Bits.of_N 32 4294906749) in let i := #(Bits.of_N 32 4294942216) in struct complex {re:=r;im:=i} }}
| 3, 8 => {{let r := #(Bits.of_N 32 4294901760) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
| 3, 9 => {{let r := #(Bits.of_N 32 4294906749) in let i := #(Bits.of_N 32 25079) in struct complex {re:=r;im:=i} }}
| 3, 10 => {{let r := #(Bits.of_N 32 4294920955) in let i := #(Bits.of_N 32 46340) in struct complex {re:=r;im:=i} }}
| 3, 11 => {{let r := #(Bits.of_N 32 4294942216) in let i := #(Bits.of_N 32 60547) in struct complex {re:=r;im:=i} }}
| 3, 12 => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 65536) in struct complex {re:=r;im:=i} }}
| 3, 13 => {{let r := #(Bits.of_N 32 25079) in let i := #(Bits.of_N 32 60547) in struct complex {re:=r;im:=i} }}
| 3, 14 => {{let r := #(Bits.of_N 32 46340) in let i := #(Bits.of_N 32 46340) in struct complex {re:=r;im:=i} }}
| 3, 15 => {{let r := #(Bits.of_N 32 60547) in let i := #(Bits.of_N 32 25079) in struct complex {re:=r;im:=i} }}
| _, _ => {{let r := #(Bits.of_N 32 0) in let i := #(Bits.of_N 32 0) in struct complex {re:=r;im:=i} }}
end.
Definition twist (stage:nat) (index:nat) : uaction reg_t ext_fn_t :=
twiddle stage index
.
Definition stage_ft (stage:nat) : UInternalFunction reg_t ext_fn_t :=
let ainit := array_init (reg_t:= reg_t) (struct_t complex) points in
{{ fun stage_ft
(stage_in: array_t {| array_type := struct_t complex ; array_len:= points|}) :
array_t {| array_type := struct_t complex ; array_len:= points|} =>
let zero := #(Bits.of_nat 32 0) in
let zerocomplex := struct complex {re:= zero; im := zero} in
let interm := ainit(zerocomplex) in
`List.fold_left (fun acc el =>
{{`acc`;
let y := bfly2(`static_slice
(struct_t complex)
({{zerocomplex}})
2
{{stage_in}}
(2*el)`,
`twist stage el`)
in
(set interm :=
`static_update
{{interm}}
(el*2)
(static_access {{y}} 0)`);
set interm :=
`static_update
{{interm}}
(el*2 + 1)
(static_access {{y}} 1)`}})
(seq 0 (points/2))
{{ pass }}`;
let out := ainit(zerocomplex) in
`List.fold_left (fun acc el =>
{{
(set out :=
`static_update
{{ out }}
el
(static_access {{interm}} (permute el points))`); `acc`}})
(seq 0 points)
{{ pass }}`;
out
}}.
Inductive rule_name_t := tb .
Definition R r :=
match r with
| localSt => array_t {| array_type := struct_t complex ; array_len:= points|}
end.
Import VectNotations.
Definition unpack_complex n :=
@value_of_bits (struct_t complex) (Bits.of_N _ n).
Import VectNotations.
Definition r idx : R idx :=
match idx with
(*| localSt => value_of_bits (Bits.of_N (64*points) 4039432 : bits_t (64*points))*)
| localSt =>
vect_map unpack_complex
[ 4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432;
4039543289543432]%N%vect
end.
Definition testbench_fir : scheduler :=
tb |> done.
Definition cr := ContextEnv.(create) r.
Definition _tb : uaction reg_t ext_fn_t:=
{{
let result :=
`List.fold_left (fun acc el =>
let fft := stage_ft el in {{fft(`acc`)}}) (seq 0 log_points) {{read0(localSt)}}` in
write0(localSt, result)
}}.
(* Typechecking *)
Definition rules :=
tc_rules R Sigma
(fun r => match r with
| tb => _tb
end).
Definition external (r: rule_name_t) := false.
Definition package :=
{| ip_koika := {| koika_reg_types := R;
koika_reg_init := r;
koika_ext_fn_types := Sigma;
koika_rules := rules;
koika_rule_external := external;
koika_scheduler := testbench_fir;
koika_module_name := "fft" |};
ip_sim := {| sp_ext_fn_specs := empty_ext_fn_props;
sp_prelude := None |};
ip_verilog := {| vp_ext_fn_specs := empty_ext_fn_props |} |}.
End FFT.
Definition prog := Interop.Backends.register package.
Extraction "fft.ml" prog.