microsoft / knossos-ksc

Compiler with automatic differentiation
Other
45 stars 10 forks source link

map doesn't work over tuple types #847

Closed toelli-msft closed 3 years ago

toelli-msft commented 3 years ago

Consider the following functio

(def map_tuple (Vec Float) (x : Vec (Tuple Float Float))
     (map (lam (t : Tuple Float Float)
               (let ((t1 t2) t) (add t1 t2)))
          x))

ksc currently generates C++ like

typedef tensor<1, double> ty$map_tuple$aT1$dff$b;
ty$map_tuple$aT1$dff$b map_tuple$aT1$dff$b(ks::allocator * $alloc, tensor<1, tuple<double,double>> x) {
  /* Lam */auto c$0 = [=](ks::allocator * $alloc, double targ1, double targ2) {
    tuple<double,double> t = ks::make_tuple(targ1,targ2);
    auto [t1, t2] = t;
    double c$1 = add$aff($alloc, t1, t2);
    return (c$1);
  };
  tensor<1, double> c$2 = map($alloc, c$0, x);
  return (c$2);
}

Because of our Cgen special rules that automatically unpack functions of tuple arguments the lam is compiled to auto c$0 = [=](ks::allocator * $alloc, double targ1, double targ2) which has too many arguments! Oops.

What's the right solution? Maybe to not unpack tuple arguments of lams? I think that is OK since I think that lams and defs are not actually interchangable. Needs more thought/discussion.

toelli-msft commented 3 years ago

What's the right solution? Maybe to not unpack tuple arguments of lams?

Hmm, no that doesn't work, because build takes a lam and does the unpacking of the tuple itself. Implementing my suggestion will break build. I suppose we have two options

  1. Implement my suggestion and change build to be compatible with it
  2. Change the implementation of map to unpack a tuple argument to the mapped function.
  3. Change Cgen to detect when it is generating a map (or suffwdpass_map) and generate special code that doesn't unpack its argument tuple.

2 is probably less invasive than 1 but I may need a C++ expert. 3 is least invasive but most special-cased.

toelli-msft commented 3 years ago

I did 2 https://github.com/microsoft/knossos-ksc/pull/848