1 ! Copyright (C) 2019 HMC Clinic.
2 ! See http://factorcode.org/license.txt for BSD license.
4 USING: accessors alien alien.c-types alien.data arrays combinators
5 grouping kernel math math.functions ranges math.vectors
6 math.vectors.simd multi-methods parser prettyprint.custom sequences sequences.extras
7 sequences.private specialized-arrays typed ;
9 QUALIFIED-WITH: alien.c-types c
10 SPECIALIZED-ARRAY: c:float
11 SPECIALIZED-ARRAY: float-4
14 ! Tensor class definition
20 ERROR: non-positive-shape-error shape ;
21 ERROR: shape-mismatch-error shape1 shape2 ;
22 ERROR: non-uniform-seq-error seq ;
23 ERROR: dimension-mismatch-error tensor-dim index-dim ;
27 ! Check that the shape has only positive values
28 : check-shape ( shape -- shape )
29 dup [ 1 < ] map-find drop [ non-positive-shape-error ] when ;
31 ! Construct a tensor of zeros
32 : <tensor> ( shape seq -- tensor )
35 ! Creates a freshly-allocated float-array with the desired c-type values
36 : >float-array ( seq -- float-array )
39 : repetition ( shape const -- tensor )
40 [ check-shape dup product ] dip <repetition>
41 >float-array <tensor> ;
45 ! Construct a tensor of zeros
46 : zeros ( shape -- tensor )
49 ! Construct a tensor of ones
50 : ones ( shape -- tensor )
53 ! Construct a one-dimensional tensor with values start, start+step,
54 ! ..., stop (inclusive)
55 : arange ( a b step -- tensor )
56 <range> [ length >fixnum 1array ] keep >float-array <tensor> ;
58 ! Construct a tensor with vec { 0 1 2 ... } and reshape to the desired shape
59 : naturals ( shape -- tensor )
60 check-shape [ ] [ product [0..b) >float-array ] bi <tensor> ;
62 ! Construct a tensor without initializing its values
63 : (tensor) ( shape -- tensor )
64 dup product (float-array) <tensor> ;
68 : check-reshape ( shape1 shape2 -- shape1 shape2 )
69 2dup [ product ] bi@ = [ shape-mismatch-error ] unless ;
73 ! Reshape the tensor to conform to the new shape
74 : reshape ( tensor shape -- tensor )
75 [ dup shape>> ] [ check-shape ] bi* check-reshape nip >>shape ;
77 ! Flatten the tensor so that it is only one-dimensional
78 : flatten ( tensor -- tensor )
80 product { } 1sequence >>shape ;
82 ! outputs the number of dimensions of a tensor
83 : dims ( tensor -- n )
86 ! Turn into Factor ND array form
87 ! Source: shaped-array>array
88 TYPED: tensor>array ( tensor: tensor -- seq: array )
89 [ vec>> >array ] [ shape>> ] bi
90 [ rest-slice reverse [ group ] each ] unless-empty ;
93 ! recursively finds shape of nested array
94 ! assumes properly shaped array (all sub-arrays are same size)
95 :: find-shape ( seq shape -- shape' )
96 seq empty? [ { 0 } ] [
97 ! add length of seq element to shape
98 shape seq length 1array append :> shape'
99 ! base case: check if the first element is a seq
102 ! is a sequence: recurse on 1st element
103 [ 1st shape' find-shape ]
104 ! not a sequence: return shape'
109 ! turns a nested array into a tensor
110 :: >tensor ( seq -- tensor )
112 seq { } find-shape :> shape
117 ] each-integer :> flatseq
118 ! check that the size is good
119 shape product flatseq length =
120 [ seq non-uniform-seq-error ] unless
122 shape flatseq >float-array <tensor> ;
124 SYNTAX: t{ \ } [ >tensor ] parse-literal ;
127 syntax:M: tensor pprint-delims drop \ t{ \ } ;
128 syntax:M: tensor >pprint-sequence tensor>array ;
129 syntax:M: tensor pprint* pprint-object ;
133 ! turns a shape into a list of things by which to multiply
134 ! indices to get a full index (e.g. { 2 3 4 } -> { 12 4 1 })
135 : ind-mults ( shape -- seq )
136 <reversed> 1 swap [ swap [ * ] keep ] map nip reverse ;
138 ! turns a num/seq index & tensor into num index & tensor
139 ! also throws a dimension mismatch if seq & tens shape>> arent the same len
140 : num-index ( n/seq tensor -- n tensor )
141 ! check form of index (num or seq)
143 [ ! if array, first check if it's a valid index
144 2dup [ shape>> length ] dip length 2dup =
145 [ dimension-mismatch-error ] unless 2drop
147 [ dup shape>> ind-mults ] dip [ * ] 2map-sum
153 ! Sequence protocol implementation
154 syntax:M: tensor clone [ shape>> clone ] [ vec>> clone ] bi <tensor> ;
156 syntax:M: tensor length vec>> length ;
158 syntax:M: tensor nth num-index vec>> nth ;
160 syntax:M: tensor nth-unsafe num-index vec>> nth-unsafe ;
162 syntax:M: tensor set-nth num-index vec>> set-nth ;
164 syntax:M: tensor set-nth-unsafe num-index vec>> set-nth-unsafe ;
166 syntax:M: tensor new-sequence
167 ! Check if the old and new tensors are the same size
168 shape>> 2dup product =
169 ! If so preserve the shape, otherwise create a 1D tensor
170 [ nip (tensor) ] [ drop 1array (tensor) ] if ;
172 syntax:M: tensor like
173 ! If the original sequence is already a tensor, we are done
177 [ dup [ length 1array ] dip <tensor> ] dip
181 2dup [ length ] bi@ = [ shape>> reshape ] [ drop ] if
184 syntax:M: tensor clone-like
185 ! If the original sequence is already a tensor, we just need to clone it
189 2dup [ length ] bi@ = [ shape>> reshape ] [ drop ] if
192 INSTANCE: tensor sequence
197 :: make-subseq ( arr start len -- arr )
199 c:float heap-size start *
200 ! Compute the starting pointer
201 arr underlying>> <displaced-alien>
202 ! Push length and type to create the new array
203 len c:float <c-direct-array> ; inline
205 : check-bop-shape ( shape1 shape2 -- shape )
206 2dup = [ shape-mismatch-error ] unless drop ;
208 ! Apply the binary operator bop to combine the tensors
209 TYPED:: t-bop ( tensor1: tensor tensor2: tensor quot: ( x y -- z ) -- tensor: tensor )
210 tensor1 shape>> tensor2 shape>> check-bop-shape
211 tensor1 vec>> tensor2 vec>> quot 2map <tensor> ; inline
213 ! Create an array of 4-element SIMD arrays for processing floats
214 : simd-for-bop ( array -- simd-array rest-slice/f )
215 dup length dup 4 mod [ drop f ] [ - cut-slice ] if-zero
216 [ float-4 cast-array ] dip ; inline
218 ! Create an array of 4-element SIMD arrays for processing floats
219 ! Tensor class definition
221 { first-slice float-array }
222 { simd-slice float-4-array }
223 { end-slice float-array } ;
225 :: (simd-slice) ( arr start len -- arr/f )
226 len [ float-array{ } ] [ drop arr start len make-subseq ] if-zero ; inline
228 :: <simd-slice> ( arr start -- simd-slice )
229 ! Compute the beginning
230 arr 0 start (simd-slice)
231 ! Compute the SIMD part
232 arr length start - :> len
234 arr start len end - (simd-slice) float-4 cast-array
236 arr dup length end - end (simd-slice)
237 simd-slice boa ; inline
239 ! Apply the binary operators simd-quot and quot to quickly combine the tensors
240 :: t-bop-simd ( tensor1 tensor2 simd-quot: ( x y -- z ) quot: ( x y -- z ) -- tensor )
241 tensor1 shape>> tensor2 shape>> check-bop-shape
242 tensor1 vec>> tensor2 vec>>
243 dup length (float-array) dup :> vec3
244 [ simd-for-bop ] tri@ :> ( simd1 rest1 simd2 rest2 simd3 rest3 )
245 simd1 simd2 simd-quot simd3 2map-into
246 rest1 rest2 quot rest3 2map-into
247 vec3 <tensor> ; inline
249 ! Apply the operation to the tensor
250 TYPED:: t-uop ( tensor: tensor quot: ( x -- y ) -- tensor: tensor )
251 tensor vec>> quot map [ tensor shape>> ] dip <tensor> ; inline
253 ! Apply the binary operators simd-quot and quot to quickly combine a tensor and
255 :: t-uop-simd ( tensor n simd-quot: ( x y -- z ) quot: ( x y -- z ) -- tensor )
256 tensor dup [ shape>> ] [ vec>> ] bi*
257 dup length (float-array) dup :> vec2
258 [ simd-for-bop ] bi@ :> ( simd1 rest1 simd2 rest2 )
259 simd1 n n n n float-4-boa simd-quot curry simd2 map-into
260 rest1 n quot curry rest2 map-into
261 vec2 <tensor> ; inline
265 ! Add a tensor to either another tensor or a scalar
266 multi-methods:GENERIC: t+ ( x y -- tensor )
267 METHOD: t+ { tensor tensor } [ v+ ] [ + ] t-bop-simd ;
268 METHOD: t+ { tensor number } >float [ v+ ] [ + ] t-uop-simd ;
269 METHOD: t+ { number tensor } swap >float [ swap v+ ] [ swap + ] t-uop-simd ;
271 ! Subtraction between two tensors or a tensor and a scalar
272 multi-methods:GENERIC: t- ( x y -- tensor )
273 METHOD: t- { tensor tensor } [ v- ] [ - ] t-bop-simd ;
274 METHOD: t- { tensor number } >float [ v- ] [ - ] t-uop-simd ;
275 METHOD: t- { number tensor } swap >float [ swap v- ] [ swap - ] t-uop-simd ;
277 ! Multiply a tensor with either another tensor or a scalar
278 multi-methods:GENERIC: t* ( x y -- tensor )
279 METHOD: t* { tensor tensor } [ v* ] [ * ] t-bop-simd ;
280 METHOD: t* { tensor number } >float [ v* ] [ * ] t-uop-simd ;
281 METHOD: t* { number tensor } swap >float [ swap v* ] [ swap * ] t-uop-simd ;
283 ! Divide two tensors or a tensor and a scalar
284 multi-methods:GENERIC: t/ ( x y -- tensor )
285 METHOD: t/ { tensor tensor } [ v/ ] [ / ] t-bop-simd ;
286 METHOD: t/ { tensor number } >float [ v/ ] [ / ] t-uop-simd ;
287 METHOD: t/ { number tensor } swap >float [ swap v/ ] [ swap / ] t-uop-simd ;
289 ! Mod two tensors or a tensor and a scalar
290 multi-methods:GENERIC: t% ( x y -- tensor )
291 METHOD: t% { tensor tensor } [ mod ] t-bop ;
292 METHOD: t% { tensor number } >float [ mod ] curry t-uop ;
293 METHOD: t% { number tensor } [ >float ] dip [ mod ] with t-uop ;
295 ! Sum together all elements in the tensor
296 syntax:M: tensor sum vec>> 0 <simd-slice>
297 [ simd-slice>> 0 [ sum + ] reduce ]
298 [ end-slice>> sum ] bi + ;
302 ! Also converts all elements of the sequence to tensors
303 :: check-concat-shape ( seq -- seq )
304 ! Compute the bottom shape of the first element in the sequence
305 seq first { } >tensor dup :> empty-tensor
306 like shape>> dup :> first-shape rest :> rest-shape
308 ! Compute the bottom shape of this element
309 empty-tensor like dup shape>> rest
310 ! Compare; if they are different, throw an error
311 rest-shape = [ shape>> first-shape swap shape-mismatch-error ] unless
314 ! Also converts all elements of the sequence to tensors
315 :: check-stack-shape ( seq -- seq )
316 ! Compute the bottom shape of the first element in the sequence
317 seq first { } >tensor dup :> empty-tensor
318 like shape>> :> first-shape
320 ! Compute the bottom shape of this element
321 empty-tensor like dup shape>>
322 ! Compare; if they are different, throw an error
323 first-shape = [ shape>> first-shape swap shape-mismatch-error ] unless
326 ! Also converts all elements of the sequence to tensors
327 :: check-hstack-shape ( seq -- seq )
328 ! Compute the top shape of the first element in the sequence
329 seq first { } >tensor dup :> empty-tensor
330 like shape>> dup :> first-shape but-last :> but-last-shape
332 ! Compute the top shape of this element
333 empty-tensor like dup shape>> but-last
334 ! Compare; if they are different, throw an error
335 but-last-shape = [ shape>> first-shape swap shape-mismatch-error ] unless
338 : final-hstack-shape ( seq -- shape )
340 dup first shape>> but-last swap
341 ! Compute the last part of the shape
342 [ shape>> last ] map sum 1array append ;
344 ! Returns an guide for hstacking where the index corresponds to the postion
345 ! in the last dimension of the resulting tensor, and the elements are
346 ! { which tensor, len of tensor, index }
347 :: hstack-guide ( seq -- guide )
348 ! Compute the list of last shape parts
349 seq [ shape>> last ] map :> last-dims
350 ! Curr tensor and index in tensor
352 last-dims sum [0..b) [
353 drop :> old-t-ind :> last-dims-i
354 last-dims-i last-dims nth
356 ! If we need to move onto the next tensor
357 [ last-dims-i 1 + 0 ]
358 ! Otherwise, stay with the current tensor
359 [ drop last-dims-i old-t-ind ] if-zero
360 2dup [ dup last-dims nth ] dip 3array
364 ! Given a sequence of tensors, stack them across the last dimension
365 :: hstack-unsafe ( tseq -- tensor )
366 ! Create the final tensor
367 tseq final-hstack-shape (tensor)
368 ! Compute the guide information
369 tseq hstack-guide dup length :> repeat :> guide
372 ! First get the correct tensor
373 i repeat /mod guide nth
375 ! Now find the correct value within that tensor
376 [ [ second ] [ third ] bi -rot * + ] dip nth
379 ! Also converts all elements of the sequence to tensors
380 :: check-vstack-shape ( seq -- seq )
381 ! Compute the shape of the first sequence
382 seq first { } >tensor dup :> empty-tensor
383 like shape>> dup :> first-shape
384 ! Compute the index of the dimension to be stacked across
387 ! Convert this element to a tensor
388 empty-tensor like dup
390 shape>> first-shape [ = ] 2map
392 ! If the shapes differ in anything except the second-to-last dimension
393 ! this sequence cannot be vstacked
394 t [ = ] reduce [ shape>> first-shape swap shape-mismatch-error ] unless
397 ! Compute the shape after the vstack has been completed
398 :: final-vstack-shape ( seq -- shape )
399 ! Compute the new second-to-last dimension
400 seq first dims 2 - :> vdim
401 seq 0 [ shape>> vdim swap nth + ] reduce
402 ! Combine it to create the new shape
403 seq first shape>> clone :> new-shape
404 vdim new-shape set-nth
407 ! Combine the second-to-last and last dimensions of each tensor for stacking
408 :: reshape-for-vstack ( seq -- seq )
409 seq first dims 2 - :> vdim
411 dup shape>> vdim cut product 1array append >>shape
417 ! Concatenation operations
418 ! Concatenate across the last dimension
419 : t-concat ( seq -- tensor )
421 ! Compute the final shape
423 ! Compute the first dimension
424 [ 0 [ shape>> first + ] reduce 1array ]
425 ! Compute the other dimensions
426 [ first shape>> rest ] bi append
428 ! Concatenate all of the float-arrays
429 [ [ vec>> ] map concat ] bi <tensor> ;
431 : stack ( seq -- tensor )
433 ! Compute the new shape
434 [ [ length 1array ] [ first shape>> ] bi append ]
435 ! Concatenate all of the tensors
436 [ [ vec>> ] map concat ] bi <tensor> ;
438 : hstack ( seq -- tensor )
439 ! Check shape and convert everything to tensors
440 check-hstack-shape hstack-unsafe ;
442 : vstack ( seq -- tensor )
443 ! Check shape and convert everything to tensors
445 ! Find the final shape
446 [ final-vstack-shape ]
447 ! Reshape each of the tensors and stack
448 [ reshape-for-vstack hstack-unsafe ] bi
449 ! Finally reshape and return
454 ! Check that the tensor has an acceptable shape for matrix multiplication
455 : check-matmul-shape ( tensor1 tensor2 -- )
456 [let [ shape>> ] bi@ :> shape2 :> shape1
457 ! Check that the matrices can be multiplied
458 shape1 last shape2 [ length 2 - ] keep nth =
459 ! Check that the other dimensions are equal
460 shape1 2 head* shape2 2 head* = and
461 ! If either is false, raise an error
462 [ shape1 shape2 shape-mismatch-error ] unless ] ;
464 ! Slice out a row from the array
465 : row ( arr n i p -- slice )
466 ! Compute the starting index
468 ! Compute the ending index
473 ! much quicker transpose for 2d tensors
474 TYPED:: 2d-transpose ( tensor: tensor -- tensor': tensor )
475 tensor shape>> :> old-shape
477 old-shape first2 :> ( s1 s2 )
478 ! loop through new tensor
479 old-shape reverse dup product <iota> [
480 ! find y*b val in original tensor
482 ! find x val in original tensor
483 [ s2 /mod ] dip + nip
484 ! get that index in original tensor
486 ] float-array{ } map-as <tensor> ;
488 ! Perform matrix multiplication muliplying an
489 ! mxn matrix with a nxp matrix
490 TYPED:: 2d-matmul ( vec1: float-array vec2: float-array res: float-array
491 m: fixnum n: fixnum p: fixnum -- )
492 ! For each element in the range, we want to compute the dot product of the
493 ! corresponding row and column
494 ! Transpose vec2 so that we are doing row * row (as opposed to row * col)
495 { n p } vec2 <tensor> 2d-transpose vec>> :> vec2
500 vec1 in n make-subseq
503 vec2 j n * n make-subseq
505 ip j + res set-nth-unsafe
510 ! Perform matrix multiplication muliplying an
511 ! mxn matrix with a nxp matrix
512 TYPED:: 2d-matmul-mixed ( vec1: float-array vec2: float-array res: float-array
513 m: fixnum n: fixnum p: fixnum start: fixnum -- )
514 ! For each element in the range, we want to compute the dot product of the
515 ! corresponding row and column
516 ! Transpose vec2 so that we are doing row * row (as opposed to row * col)
517 { n p } vec2 <tensor> 2d-transpose vec>> :> vec2
519 ! Compute the location in the float-array each 2D matrix will start at
520 start m n * * :> start1
521 start n p * * :> start2
525 4 4 in start1 + 4 mod - swap mod :> in4m
527 vec1 in n make-subseq :> sub1
528 sub1 in4m <simd-slice> :> slice1
531 4 4 jn 4 mod - swap mod :> jn4m
532 vec2 jn n make-subseq
534 jn4m <simd-slice> slice1 swap
535 2dup [ first-slice>> ] bi@ 0.0 [ * + ] 2reduce
536 [ 2dup [ simd-slice>> ] bi@ ] dip [ vdot + ] 2reduce
537 [ [ end-slice>> ] bi@ ] dip [ * + ] 2reduce
542 ip j + res set-nth-unsafe
546 ! ! Perform matrix multiplication muliplying an
547 ! mxn matrix with a nxp matrix
548 ! Should only be called when n is a multiple of 4
549 TYPED:: 2d-matmul-simd ( vec1: float-array vec2: float-array
551 m: fixnum n: fixnum p: fixnum -- )
552 ! For each element in the range, we want to compute the dot product of the
553 ! corresponding row and column
554 ! Transpose vec2 so that we are doing row * row (as opposed to row * col)
555 { n p } vec2 <tensor> 2d-transpose vec>> :> vec2
560 vec1 in n make-subseq float-4 cast-array
563 vec2 j n * n make-subseq float-4 cast-array
564 0.0 [ vdot + ] 2reduce
565 ip j + res set-nth-unsafe
573 ! Perform matrix multiplication muliplying an
574 ! ...xmxn matrix with a ...xnxp matrix
575 TYPED:: matmul ( tensor1: tensor tensor2: tensor -- tensor3: tensor )
576 ! First check the shape
577 tensor1 tensor2 check-matmul-shape
579 ! Now save all of the sizes
580 tensor1 shape>> unclip-last-slice :> n
581 unclip-last-slice :> m :> top-shape
582 tensor2 shape>> last :> p
583 top-shape product :> top-prod
585 ! Create the shape of the resulting tensor
586 top-shape { m p } append
588 ! Now create the new float array to store the underlying result
589 dup product (float-array) :> vec3
591 ! Now update the tensor3 to contain the multiplied matricies
595 ! Compute vec1 using direct C arrays
596 tensor1 vec>> m n * i * m n * make-subseq
598 ! Compute vec2 and start2
599 tensor2 vec>> n p * i * n p * make-subseq
602 vec3 m p * i * m p * make-subseq
603 ! Push m, n, and p and multiply the arrays
605 { { [ n 4 mod 0 = ] [ 2d-matmul-simd ] }
606 { [ n 4 < ] [ 2d-matmul ] }
607 [ i 2d-matmul-mixed ]
613 ! Transpose an n-dimensional tensor by flipping the axes
614 TYPED:: transpose ( tensor: tensor -- tensor': tensor )
615 tensor shape>> length 2 =
616 [ tensor 2d-transpose ]
617 [ tensor shape>> :> old-shape
619 old-shape reverse :> new-shape
620 old-shape ind-mults :> mults
621 ! loop through new tensor
622 new-shape dup product <iota> [
623 ! find index in original tensor
624 old-shape mults [ [ /mod ] dip * ] 2map-sum nip
625 ! get that index in original tensor
627 ] float-array{ } map-as <tensor>