basis/math/vectors/specialization/specialization.factor

   1 ! Copyright (C) 2009 Slava Pestov.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: alien.c-types words kernel make sequences effects
   4 kernel.private accessors combinators math math.intervals
   5 math.vectors namespaces assocs fry splitting classes.algebra
   6 generalizations locals compiler.tree.propagation.info ;
   7 IN: math.vectors.specialization
   8
   9 SYMBOLS: -> +vector+ +scalar+ +nonnegative+ ;
  10
  11 : signature-for-schema ( array-type elt-type schema -- signature )
  12     [
  13         {
  14             { +vector+ [ drop ] }
  15             { +scalar+ [ nip ] }
  16             { +nonnegative+ [ nip ] }
  17         } case
  18     ] with with map ;
  19
  20 : (specialize-vector-word) ( word array-type elt-type schema -- word' )
  21     signature-for-schema
  22     [ [ name>> ] [ [ name>> ] map "," join ] bi* "=>" glue f <word> ]
  23     [ [ , \ declare , def>> % ] [ ] make ]
  24     [ drop stack-effect ]
  25     2tri
  26     [ define-declared ] [ 2drop ] 3bi ;
  27
  28 : output-infos ( array-type elt-type schema -- value-infos )
  29     [
  30         {
  31             { +vector+ [ drop <class-info> ] }
  32             { +scalar+ [ nip <class-info> ] }
  33             { +nonnegative+ [ nip real class-and [0,inf] <class/interval-info> ] }
  34         } case
  35     ] with with map ;
  36
  37 : record-output-signature ( word array-type elt-type schema -- word )
  38     output-infos
  39     [ drop ]
  40     [ drop ]
  41     [ [ stack-effect in>> length '[ _ ndrop ] ] dip append ] 2tri
  42     "outputs" set-word-prop ;
  43
  44 CONSTANT: vector-words
  45 H{
  46     { [v-] { +vector+ +vector+ -> +vector+ } }
  47     { distance { +vector+ +vector+ -> +nonnegative+ } }
  48     { n*v { +scalar+ +vector+ -> +vector+ } }
  49     { n+v { +scalar+ +vector+ -> +vector+ } }
  50     { n-v { +scalar+ +vector+ -> +vector+ } }
  51     { n/v { +scalar+ +vector+ -> +vector+ } }
  52     { norm { +vector+ -> +nonnegative+ } }
  53     { norm-sq { +vector+ -> +nonnegative+ } }
  54     { normalize { +vector+ -> +vector+ } }
  55     { v* { +vector+ +vector+ -> +vector+ } }
  56     { v*n { +vector+ +scalar+ -> +vector+ } }
  57     { v+ { +vector+ +vector+ -> +vector+ } }
  58     { v+n { +vector+ +scalar+ -> +vector+ } }
  59     { v- { +vector+ +vector+ -> +vector+ } }
  60     { v-n { +vector+ +scalar+ -> +vector+ } }
  61     { v. { +vector+ +vector+ -> +scalar+ } }
  62     { v/ { +vector+ +vector+ -> +vector+ } }
  63     { v/n { +vector+ +scalar+ -> +vector+ } }
  64     { vceiling { +vector+ -> +vector+ } }
  65     { vfloor { +vector+ -> +vector+ } }
  66     { vmax { +vector+ +vector+ -> +vector+ } }
  67     { vmin { +vector+ +vector+ -> +vector+ } }
  68     { vneg { +vector+ -> +vector+ } }
  69     { vtruncate { +vector+ -> +vector+ } }
  70     { sum { +vector+ -> +scalar+ } }
  71 }
  72
  73 PREDICATE: vector-word < word vector-words key? ;
  74
  75 : specializations ( word -- assoc )
  76     dup "specializations" word-prop
  77     [ ] [ V{ } clone [ "specializations" set-word-prop ] keep ] ?if ;
  78
  79 M: vector-word subwords specializations values [ word? ] filter ;
  80
  81 : add-specialization ( new-word signature word -- )
  82     specializations set-at ;
  83
  84 : word-schema ( word -- schema ) vector-words at ;
  85
  86 : inputs ( schema -- seq ) { -> } split first ;
  87
  88 : outputs ( schema -- seq ) { -> } split second ;
  89
  90 : loop-vector-op ( word array-type elt-type -- word' )
  91     pick word-schema
  92     [ inputs (specialize-vector-word) ]
  93     [ outputs record-output-signature ] 3bi ;
  94
  95 :: specialize-vector-word ( word array-type elt-type simd -- word/quot' )
  96     word simd key? [ word simd at ] [ word array-type elt-type loop-vector-op ] if ;
  97
  98 :: input-signature ( word array-type elt-type -- signature )
  99     array-type elt-type word word-schema inputs signature-for-schema ;
 100
 101 :: specialize-vector-words ( array-type elt-type simd -- )
 102     elt-type number class<= [
 103         vector-words keys [
 104             [ array-type elt-type simd specialize-vector-word ]
 105             [ array-type elt-type input-signature ]
 106             [ ]
 107             tri add-specialization
 108         ] each
 109     ] when ;
 110
 111 : find-specialization ( classes word -- word/f )
 112     specializations
 113     [ first [ class<= ] 2all? ] with find
 114     swap [ second ] when ;
 115
 116 : vector-word-custom-inlining ( #call -- word/f )
 117     [ in-d>> [ value-info class>> ] map ] [ word>> ] bi
 118     find-specialization ;
 119
 120 vector-words keys [
 121     [ vector-word-custom-inlining ]
 122     "custom-inlining" set-word-prop
 123 ] each