tensor_file.bend

stelliferous@0.0.2.0 · 422 lines · 79 definitions

1# Raw little-endian FP32 files and NumPy .npy files through the official File effects.
2import Base
3import stelliferous@0.0.2.0/storage_buffer.bend as Buffers
4import stelliferous@0.0.2.0/tensor.bend as Tensors
5import stelliferous@0.0.2.0/tensor_layout.bend as Layout
6
7def Tensor() -> Type:
8  Tensors.Tensor<F32>
9
10def Read() -> Type:
11  File & Result<&1,&1,U32 & String,List<&2,U32>>
12
13def Written() -> Type:
14  File & Result<&1,&1,U32 & String,Unit>
15
16# Values per File effect.
17def chunk() -> U32:
18  16384
19
20# An FP32 value and its 32 bits are the same word.
21
22def from_bits(word: U32) -> F32:
23  U32{bits} = word
24  F32{bits}
25
26def to_bits(value: F32) -> U32:
27  F32{bits} = value
28  U32{bits}
29
30# Reading: bytes gather into words, least significant first. The values of one
31# read are written into the buffer before the next read, so no list outgrows a read.
32
33type Words is Data:
34  Words{values: List<&2,F32>,count: U32,word: U32,weight: U32}
35
36# The buffer of the tensor, the values found so far, and a word in progress.
37type Target is Type:
38  Target{buffer: Buffers.Buffer<F32>,found: U32,word: U32,weight: U32}
39
40type Reading is Type:
41  More{file: File,target: Target}
42  End{file: File,target: Target,complete: Bool}
43
44type Outcome is Type:
45  Outcome{target: Target,complete: Bool}
46
47def gather(full: Bool,values: List<&2,F32>,+count: U32,+word: U32,+weight: U32) -> Words:
48  match full:
49    case False{}: Words{values,count,word,(weight * 256 : U32)}
50    case True{}: Words{from_bits(word) <> values,(count + 1 : U32),0,1}
51
52def push(+byte: U32,words: Words) -> Words:
53  Words{values,+count,+word,+weight} = words
54  gather(U32.is_eq(weight,16777216),values,count,(word + byte * weight : U32),weight)
55
56def trickle(bytes: List<&2,U32>,words: Words) -> Words:
57  match bytes:
58    case []: words
59    case byte <> rest: trickle(rest,push(byte,words))
60
61# Four bytes at a time while the bytes are whole words; the values are newest first.
62def quads(bytes: List<&2,U32>,values: List<&2,F32>,+count: U32) -> Words:
63  match bytes:
64    case Con{first,Con{second,Con{third,Con{fourth,rest}}}}:
65      quads(rest,from_bits((first + second * 256 + third * 65536 + fourth * 16777216 : U32)) <> values,(count + 1 : U32))
66    case rest: trickle(rest,Words{values,count,0,1})
67
68def words(whole: Bool,bytes: List<&2,U32>,+word: U32,+weight: U32) -> Words:
69  match whole:
70    case True{}: quads(bytes,[],0)
71    case False{}: trickle(bytes,Words{[],0,word,weight})
72
73# Values beyond the count of the shape are counted and not written.
74def placed(read: Words,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Target:
75  Words{values,+added,+word,+weight} = read
76  Target{Buffers.write_at(~F32,List.take(&2,F32,List.reverse(&2,F32,values),U32.to_nat((count - U32.min(found,count) : U32))),U32.min(found,count),buffer),
77    (found + added : U32),word,weight}
78
79def fill(bytes: List<&2,U32>,target: Target,+count: U32) -> Target:
80  Target{buffer,+found,+word,+weight} = target
81  placed(words(U32.is_eq(weight,1),bytes,word,weight),buffer,found,count)
82
83def absorbed(bytes: List<&2,U32>,file: File,target: Target,+count: U32) -> Reading:
84  match bytes:
85    case []: End{file,target,True{}}
86    case byte <> rest: More{file,fill(byte <> rest,target,count)}
87
88def received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,target: Target,+count: U32) -> Reading:
89  match result:
90    case Fail{error}: End{file,target,False{}}
91    case Done{bytes}: absorbed(bytes,file,target,count)
92
93def taken(read: Read(),target: Target,+count: U32) -> Reading:
94  (file,result) = read
95  received(result,file,target,count)
96
97def closed(file: File,target: Target,complete: Bool) -> IO(Outcome):
98  do IO<Outcome>:
99    File.close(file)
100    IO.pure(Outcome,Outcome{target,complete})
101
102# A file longer than its shape stops at the fuel and fails the count.
103def read_all(fuel: Nat,reading: Reading,+count: U32) -> IO(Outcome):
104  match fuel reading:
105    case 0n More{file,target}: closed(file,target,True{})
106    case 0n End{file,target,complete}: closed(file,target,complete)
107    case 1n+rest End{file,target,complete}: closed(file,target,complete)
108    case 1n+rest More{file,target}:
109      do IO<Outcome>:
110        read : Read() <- File.read_bytes(file,(chunk() * 4 : U32))
111        read_all(rest,taken(read,target,count),count)
112
113def fresh(+count: U32) -> Target:
114  Target{Tensors.buffer(~F32,0.0,count),0,0,1}
115
116def fuel(+count: U32) -> Nat:
117  U32.to_nat((count / chunk() + 2 : U32))
118
119def fault(operation: String,path: String,detail: String,shape: List<&2,U32>) -> Tensor():
120  Tensors.failed(~F32,operation,Tensors.FileFault{path,detail},shape,0.0,1)
121
122def sized(matches: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>) -> Tensor():
123  match matches:
124    case True{}: Tensors.from_storage(~F32,shape,buffer,0.0)
125    case False{}: fault(operation,path,"does not hold the values of the shape",shape)
126
127def whole(aligned: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Tensor():
128  match aligned:
129    case False{}: fault(operation,path,"is not a whole number of FP32 values",shape)
130    case True{}: sized(U32.is_eq(found,count),operation,path,shape,buffer)
131
132def aligned(operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor():
133  Target{buffer,+found,word,+weight} = target
134  whole(U32.is_eq(weight,1),operation,path,shape,buffer,found,count)
135
136def checked(complete: Bool,operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor():
137  match complete:
138    case False{}: fault(operation,path,"cannot be read",shape)
139    case True{}: aligned(operation,path,shape,target,count)
140
141def loaded(operation: String,path: String,+shape: List<&2,U32>,outcome: Outcome,+count: U32) -> Tensor():
142  Outcome{target,complete} = outcome
143  checked(complete,operation,path,shape,target,count)
144
145def opened(result: Result<&1,&1,U32 & String,File>,path: String,+shape: List<&2,U32>,+count: U32) -> IO(Tensor()):
146  match result:
147    case Fail{error}: IO.pure(Tensor(),fault("load",path,"cannot be opened",shape))
148    case Done{file}:
149      do IO<Tensor()>:
150        outcome : Outcome <- read_all(fuel(count),More{file,fresh(count)},count)
151        IO.pure(Tensor(),loaded("load",path,shape,outcome,count))
152
153def counted(count: Maybe<&2,U32>,+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
154  match count:
155    case None{}: IO.pure(Tensor(),Tensors.failed(~F32,"load",Tensors.InvalidShape{shape},shape,0.0,1))
156    case Some{+count}:
157      do IO<Tensor()>:
158        result : Result<&1,&1,U32 & String,File> <- File.open(path,"r")
159        opened(result,path,shape,count)
160
161# The file holds the values of the shape in row-major order and nothing else.
162def load(+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
163  counted(Tensors.count_of(shape),path,shape)
164
165# The NumPy .npy header: magic, version, header length, then a text such as
166# {'descr': '<f4', 'fortran_order': False, 'shape': (16, 3, 6, 6), }
167
168type Header is Data:
169  Header{shape: List<&2,U32>,data: List<&2,U32>}
170
171def codes(text: String) -> List<&2,U32>:
172  match text:
173    case SNil{}: []
174    case SCon{head,tail}: Char.to_u32(head) <> codes(tail)
175
176def magic() -> List<&2,U32>:
177  [147,78,85,77,80,89]
178
179def starts(pattern: List<&2,U32>,bytes: List<&2,U32>) -> Bool:
180  match pattern bytes:
181    case Nil{} bytes: True{}
182    case Con{first,rest} Nil{}: False{}
183    case Con{first,rest} Con{byte,more}: U32.is_eq(first,byte) && starts(rest,more)
184
185def at(bytes: List<&2,U32>,index: Nat) -> U32:
186  match bytes index:
187    case Nil{} index: 0
188    case Con{byte,rest} 0n: byte
189    case Con{byte,rest} 1n+before: at(rest,before)
190
191def contains(text: List<&2,U32>,+pattern: List<&2,U32>) -> Bool:
192  match text:
193    case Nil{}: False{}
194    case Con{+byte,+rest}: starts(pattern,byte <> rest) || contains(rest,pattern)
195
196# The text after the first occurrence of the pattern.
197def after(text: List<&2,U32>,+pattern: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
198  match text:
199    case Nil{}: None{}
200    case Con{+byte,+rest}:
201      Bool.pick(Maybe<&2,List<&2,U32>>,starts(pattern,byte <> rest),Some{List.drop(&2,U32,byte <> rest,List.length(&2,U32,pattern))},after(rest,pattern))
202
203def pushed(digits: Bool,current: U32,shape: List<&2,U32>) -> List<&2,U32>:
204  match digits:
205    case True{}: current <> shape
206    case False{}: shape
207
208def closed_shape(valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
209  match valid:
210    case True{}: Some{List.reverse(&2,U32,shape)}
211    case False{}: None{}
212
213# Extents up to the closing parenthesis: digits, commas and spaces.
214def extents(text: List<&2,U32>,+current: U32,digits: Bool,valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
215  match text:
216    case Nil{}: None{}
217    case Con{+byte,rest}:
218      match byte:
219        case 41: closed_shape(valid,pushed(digits,current,shape))
220        case 44: extents(rest,0,False{},valid,pushed(digits,current,shape))
221        case 32: extents(rest,current,digits,valid,shape)
222        case _: extents(rest,(current * 10 + byte - 48 : U32),True{},valid && U32.is_le(48,byte) && U32.is_le(byte,57),shape)
223
224def shaped(shape: Maybe<&2,List<&2,U32>>,data: List<&2,U32>) -> Maybe<&2,Header>:
225  match shape:
226    case None{}: None{}
227    case Some{shape}: Some{Header{shape,data}}
228
229def typed(supported: Bool,text: Maybe<&2,List<&2,U32>>,data: List<&2,U32>) -> Maybe<&2,Header>:
230  match supported text:
231    case False{} text: None{}
232    case True{} None{}: None{}
233    case True{} Some{text}: shaped(extents(text,0,False{},True{},[]),data)
234
235def described(+text: List<&2,U32>,+length: Nat,data: List<&2,U32>) -> Maybe<&2,Header>:
236  typed(Nat.is_eq(List.length(&2,U32,text),length) && contains(text,codes("'descr': '<f4'")) && contains(text,codes("'fortran_order': False")),
237    after(text,codes("'shape': (")),data)
238
239def region(+bytes: List<&2,U32>,+length: Nat) -> Maybe<&2,Header>:
240  described(List.take(&2,U32,bytes,length),length,List.drop(&2,U32,bytes,length))
241
242# Version 1 has a two-byte header length; later versions have four bytes.
243def versioned(major: U32,+bytes: List<&2,U32>) -> Maybe<&2,Header>:
244  match major:
245    case 1: region(List.drop(&2,U32,bytes,10n),U32.to_nat((at(bytes,8n) + at(bytes,9n) * 256 : U32)))
246    case _: region(List.drop(&2,U32,bytes,12n),U32.to_nat((at(bytes,8n) + at(bytes,9n) * 256 + at(bytes,10n) * 65536 + at(bytes,11n) * 16777216 : U32)))
247
248def marked(valid: Bool,+bytes: List<&2,U32>) -> Maybe<&2,Header>:
249  match valid:
250    case False{}: None{}
251    case True{}: versioned(at(bytes,6n),bytes)
252
253def header(+bytes: List<&2,U32>) -> Maybe<&2,Header>:
254  marked(starts(magic(),bytes),bytes)
255
256def refused(file: File,tensor: Tensor()) -> IO(Tensor()):
257  do IO<Tensor()>:
258    File.close(file)
259    IO.pure(Tensor(),tensor)
260
261def npy_counted(count: Maybe<&2,U32>,+shape: List<&2,U32>,data: List<&2,U32>,file: File,+path: String) -> IO(Tensor()):
262  match count:
263    case None{}: refused(file,Tensors.failed(~F32,"load_npy",Tensors.InvalidShape{shape},shape,0.0,1))
264    case Some{+count}:
265      do IO<Tensor()>:
266        outcome : Outcome <- read_all(fuel(count),More{file,fill(data,fresh(count),count)},count)
267        IO.pure(Tensor(),loaded("load_npy",path,shape,outcome,count))
268
269def npy_parsed(parsed: Maybe<&2,Header>,file: File,+path: String) -> IO(Tensor()):
270  match parsed:
271    case None{}: refused(file,fault("load_npy",path,"is not a little-endian FP32 .npy array in C order",[]))
272    case Some{Header{+shape,data}}: npy_counted(Tensors.count_of(shape),shape,data,file,path)
273
274def npy_received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,+path: String) -> IO(Tensor()):
275  match result:
276    case Fail{error}: refused(file,fault("load_npy",path,"cannot be read",[]))
277    case Done{bytes}: npy_parsed(header(bytes),file,path)
278
279def npy_headed(read: Read(),+path: String) -> IO(Tensor()):
280  (file,result) = read
281  npy_received(result,file,path)
282
283def npy_opened(result: Result<&1,&1,U32 & String,File>,+path: String) -> IO(Tensor()):
284  match result:
285    case Fail{error}: IO.pure(Tensor(),fault("load_npy",path,"cannot be opened",[]))
286    case Done{file}:
287      do IO<Tensor()>:
288        read : Read() <- File.read_bytes(file,(chunk() * 4 : U32))
289        npy_headed(read,path)
290
291# An FP32 array saved by NumPy, with its shape.
292def load_npy(+path: String) -> IO(Tensor()):
293  do IO<Tensor()>:
294    result : Result<&1,&1,U32 & String,File> <- File.open(path,"r")
295    npy_opened(result,path)
296
297# Writing
298
299type Writing is Type:
300  Pending{file: File,storage: Array<F32>,offset: U32,remaining: Nat}
301  Finished{file: File,complete: Bool}
302
303def bytes(+word: U32,rest: List<&2,U32>) -> List<&2,U32>:
304  (word % 256 : U32) <> (word / 256 % 256 : U32) <> (word / 65536 % 256 : U32) <> (word / 16777216 : U32) <> rest
305
306def encode(count: Nat,values: List<&2,F32>) -> List<&2,U32>:
307  match count values:
308    case 0n values: []
309    case 1n+rest Nil{}: []
310    case 1n+rest Con{value,more}: bytes(to_bits(value),encode(rest,more))
311
312# The cursor retains the array; only the current chunk becomes a list.
313def read_chunk(count: Nat,+offset: U32,storage: Array<F32>,values: List<&2,F32>) -> Array<F32> & List<&2,F32>:
314  match count:
315    case 0n: (storage,List.reverse(&2,F32,values))
316    case 1n+rest:
317      Buffers.bind_pair(Array<F32>,F32,Array<F32> & List<&2,F32>,Array.get(F32,storage,offset),storage => value =>
318        read_chunk(rest,(offset + 1 : U32),storage,value <> values))
319
320def pending(remaining: Nat,file: File,storage: Array<F32>,offset: U32) -> Writing:
321  match remaining:
322    case 0n: Finished{file,True{}}
323    case 1n+rest: Pending{file,storage,offset,1n+rest}
324
325def continued(result: Result<&1,&1,U32 & String,Unit>,file: File,storage: Array<F32>,remaining: Nat,offset: U32) -> Writing:
326  match result:
327    case Fail{error}: Finished{file,False{}}
328    case Done{unit}: pending(remaining,file,storage,offset)
329
330def after_write(written: Written(),storage: Array<F32>,remaining: Nat,offset: U32) -> Writing:
331  (file,result) = written
332  continued(result,file,storage,remaining,offset)
333
334def write_chunk(file: File,read: Array<F32> & List<&2,F32>,+count: Nat,remaining: Nat,offset: U32) -> IO(Writing):
335  (storage,values) = read
336  do IO<Writing>:
337    written : Written() <- File.write_bytes(file,encode(count,values))
338    IO.pure(Writing,after_write(written,storage,Nat.sub(remaining,count),(offset + U32.from_nat(count) : U32)))
339
340def finished(file: File,complete: Bool) -> IO(Bool):
341  do IO<Bool>:
342    File.close(file)
343    IO.pure(Bool,complete)
344
345def write_all(fuel: Nat,writing: Writing) -> IO(Bool):
346  match fuel writing:
347    case 0n Pending{file,storage,offset,remaining}: finished(file,False{})
348    case 0n Finished{file,complete}: finished(file,complete)
349    case 1n+rest Finished{file,complete}: finished(file,complete)
350    case 1n+rest Pending{file,storage,+offset,+remaining}:
351      +count = Nat.min(U32.to_nat(chunk()),remaining)
352      do IO<Bool>:
353        next : Writing <- write_chunk(file,read_chunk(count,offset,storage,[]),count,remaining,offset)
354        write_all(rest,next)
355
356def unsaved(pair: Tensor() & List<&2,U32>,operation: String,path: String,detail: String) -> Tensor():
357  (tensor,shape) = pair
358  fault(operation,path,detail,shape)
359
360def written(complete: Bool,operation: String,path: String,tensor: Tensor()) -> Tensor():
361  match complete:
362    case True{}: tensor
363    case False{}: unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be written")
364
365# The prefix is written first; a raw file has none.
366def created(result: Result<&1,&1,U32 & String,File>,operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,storage: Array<F32>,+length: Nat) -> IO(Tensor()):
367  match result:
368    case Fail{error}: IO.pure(Tensor(),unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be created"))
369    case Done{file}:
370      do IO<Tensor()>:
371        head : Written() <- File.write_bytes(file,prefix)
372        complete : Bool <- write_all(Nat.add(length,2n),after_write(head,storage,length,0))
373        IO.pure(Tensor(),written(complete,operation,path,tensor))
374
375def saving(operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,buffer: Result<&2,&1,Tensors.Problem,Buffers.Buffer<F32>>) -> IO(Tensor()):
376  match buffer:
377    case Fail{problem}: IO.pure(Tensor(),tensor)
378    case Done{Buffers.Buffer{storage,length,certificate}}:
379      do IO<Tensor()>:
380        result : Result<&1,&1,U32 & String,File> <- File.open(path,"w")
381        created(result,operation,path,tensor,prefix,storage,length)
382
383# Writes the values in row-major order and returns the tensor. A failed tensor
384# writes nothing and keeps its problem.
385def save(+path: String,tensor: Tensor()) -> IO(Tensor()):
386  Tensors.then(~Tensor(),~Tensor(),~IO(Tensor()),Layout.copy(~F32,tensor),kept => copy => saving("save",path,kept,[],Tensors.into_storage(~F32,copy)))
387
388def joined(shape: List<&2,U32>) -> String:
389  match shape:
390    case []: ""
391    case [extent]: U32.show(extent)
392    case extent <> rest: U32.show(extent) ++ ", " ++ joined(rest)
393
394# NumPy writes (), (3,) and (2, 3).
395def tuple(shape: List<&2,U32>) -> String:
396  match shape:
397    case []: "()"
398    case [extent]: "(" ++ U32.show(extent) ++ ",)"
399    case extent <> rest: "(" ++ joined(extent <> rest) ++ ")"
400
401def spaces(count: Nat) -> List<&2,U32>:
402  match count:
403    case 0n: [10]
404    case 1n+rest: 32 <> spaces(rest)
405
406# Version 1: the header ends in a newline and the data starts at a multiple of 64 bytes.
407def framed(+text: List<&2,U32>,+length: U32) -> List<&2,U32>:
408  +padding = ((64 - (11 + length) % 64) % 64 : U32)
409  +total = (length + padding + 1 : U32)
410  List.append(&2,U32,magic(),1 <> 0 <> (total % 256 : U32) <> (total / 256 : U32) <> List.append(&2,U32,text,spaces(U32.to_nat(padding))))
411
412def npy_header(shape: List<&2,U32>) -> List<&2,U32>:
413  +text = codes("{'descr': '<f4', 'fortran_order': False, 'shape': " ++ tuple(shape) ++ ", }")
414  framed(text,U32.from_nat(List.length(&2,U32,text)))
415
416def npy_saving(+path: String,pair: Tensor() & Tensor(),+shape: List<&2,U32>) -> IO(Tensor()):
417  (kept,copy) = pair
418  saving("save_npy",path,kept,npy_header(shape),Tensors.into_storage(~F32,copy))
419
420# Writes a .npy file that NumPy loads with the tensor's shape.
421def save_npy(+path: String,tensor: Tensor()) -> IO(Tensor()):
422  Tensors.then(~Tensor(),~List<&2,U32>,~IO(Tensor()),Tensors.shape(~F32,tensor),tensor => shape => npy_saving(path,Layout.copy(~F32,tensor),shape))