futhark icon indicating copy to clipboard operation
futhark copied to clipboard

Automap support.

Open zfnmxt opened this issue 2 years ago • 31 comments

In progress implementation of automap.

zfnmxt avatar Oct 20 '22 16:10 zfnmxt

Great! I have taken the liberty of merging master into this branch.

I suggest the next immediate step is to ensure that the entire test suite (including new automap test programs) type-check (futhark test -t). The internalisation can come later. This will also help reveal whether our simplifying assumptions in the design of the type system are too, well, simple.

athas avatar Oct 20 '22 16:10 athas

@athas Projections aside (which I suggest we leave until later), everything typechecks now.

zfnmxt avatar Oct 22 '22 01:10 zfnmxt

What is automap? Is it that the compiler automagically insert map as necessary?

rowanG077 avatar Oct 23 '22 14:10 rowanG077

What is automap? Is it that the compiler automagically insert map as necessary?

Yes, essentially. I wrote very lightly about it here, but it has mostly been the subject of in-person conversation: https://futhark-lang.org/blog/2022-10-03-futhark-on-arraycast.html#rank-polymorphism-in-future-futhark

athas avatar Oct 23 '22 14:10 athas

@athas Projections aside (which I suggest we leave until later), everything typechecks now.

What is difficult about projections?

athas avatar Oct 23 '22 14:10 athas

@athas Thanks! One of my professors, Sven-Bodo Scholz, is one of the prime people behind SaC and we had a course about it. Cool stuff!

rowanG077 avatar Oct 23 '22 14:10 rowanG077

@athas Projections aside (which I suggest we leave until later), everything typechecks now.

What is difficult about projections?

Nothing in particular that I'm aware of---it's just more of the same; seemed better to tack on after once we confirm the viability of the approach.

zfnmxt avatar Oct 23 '22 18:10 zfnmxt

This does not type-check:

def f (xs: []bool) (ys: []bool) = xs || ys

athas avatar Oct 24 '22 10:10 athas

My use case is expressing this in a more concise way: https://github.com/codereport/LeetCode/blob/master/0299_Problem_1.fut

athas avatar Oct 24 '22 14:10 athas

I fixed the small bool example; can you give me your automapped version of the LeetCode problem?

zfnmxt avatar Oct 25 '22 01:10 zfnmxt

Ideally this should work:

def outerprod f x y = map (f >-> flip map y) x
def bidd A = outerprod (==) (indices A) (indices A)
def xmat A = bidd A || reverse (bidd A)
def check_matrix A = xmat A == (A!=0) |> flatten |> and

But I am unsure whether we can really get away with not having any type annotations.

athas avatar Oct 25 '22 09:10 athas

With the current approach, I don't think that's possible. For example, xmat A == (A!=0) is ambiguous by itself; we know that xmat : [][]bool and A : []t, but don't know anything about the dimension of t. It could be that A: [][][]i32, which would necessitate that (xmat A ==) is mapped over (A!=0). Or it could be that A: []i32, which would necessitate (== (A !=0)) being mapped. Only the final and disambiguates things.

We only consider local information; allowing unknown type variables in the arguments that can only be resolved by type-checking subsequent parts of the expressions aren't compatible with the approach. We could, of course, take a more sophisticated approach but it would require modifications to unification.

zfnmxt avatar Oct 25 '22 14:10 zfnmxt

Remind me, what is the state of this now? Do you still want to work on it?

athas avatar Nov 24 '22 21:11 athas

Working on it right now, doing the internalization stuff we discussed. So, yes, I do want to work on it and I'm making it my priority now.

zfnmxt avatar Nov 24 '22 21:11 zfnmxt

Great to hear! Remember: the priority is making the actual map-insertion in the internaliser robust enough that we don't have to touch it again, and then focus on ever more elaborate ways of inferring the automap-annotations in the type checker.

athas avatar Nov 24 '22 21:11 athas

Yep, that's the plan. I've cleaned up all the other places where Automap touched before (a.k.a bad experiments), so right now it's only present in the type checker and internalizer (AST annotations aside).

zfnmxt avatar Nov 24 '22 21:11 zfnmxt

I see you claim all the tests are working now. The next step is to clean up the code so CI turns green. Then I will add more tests that fail, and then we will try to break our heads coming up with an overly complicated type system.

athas avatar Nov 30 '22 21:11 athas

Yep, already on it. The code is terrible right now, but at least I now know what a reasonable design looks like.

zfnmxt avatar Nov 30 '22 22:11 zfnmxt

There's something wonky going on with the short-circuiting operators. Here's the diff (the change is essentially just switching from the source IR to the core IR, since withAutoMap spits out core IR variables---which are (potentially) parameters of a map lambda---to replace the original expressions):

       -- Some functions are magical (overloaded) and we handle that here.
       case () of
         ()
           -- Short-circuiting operators are magical.
           | baseTag (qualLeaf qfname) <= maxIntrinsicTag,
-            baseString (qualLeaf qfname) == "&&",
-            [(x, _), (y, _)] <- args ->
-              internaliseExp desc $
-                E.AppExp
-                  (E.If x y (E.Literal (E.BoolValue False) mempty) mempty)
-                  (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])
+            baseString (qualLeaf qfname) == "&&" ->
+              withAutoMap ams arg_desc args $ \[[x], [y]] -> do
+                letValExp' desc
+                  =<< eIf
+                    (pure $ BasicOp $ SubExp x)
+                    (eBody [pure $ BasicOp $ SubExp y])
+                    (eBody [pure $ BasicOp $ SubExp $ Constant $ I.BoolValue $ False])

When applied to tests/eqarrays0.fut:

def main [n][m] (xs: [n]i32) (ys: [m]i32) =
  n == m && xs == (ys :> [n]i32)

the generated IR statically concludes that n is equal to m (or, at least just compares n with n):

  main_5952 (n_5974 : i64,
             m_5975 : i64,
             xs_5976 : [n_5974]i32,
             ys_5977 : [m_5975]i32)
  : {bool} = {
  let {dim_match_5978 : bool} = eq_i64(n_5974, m_5975)
  let {empty_or_match_cert_5979 : unit} =
    assert(dim_match_5978, {"Value of (core language) shape (", m_5975 : i64, ") cannot match shape of ty
pe `[", n_5974 : i64, "]i32`."}, "eqarrays0.fut:18:20-31")
  let {==_arg_5980 : [n_5974]i32} =
    #{empty_or_match_cert_5979}
    coerce([n_5974], ys_5977)
  let {dim_eq_5981 : bool} = eq_i64(n_5974, n_5974)
  let {arrays_equal_5982 : bool} =
    if  dim_eq_5981
    then {
      let {x_num_elems_5983 : i64} = mul_nw64(n_5974, 1i64)
      let {x_flat_5984 : [x_num_elems_5983]i32} =outside of the
        reshape([x_num_elems_5983], xs_5976)
      let {y_flat_5985 : [x_num_elems_5983]i32} =
        reshape([x_num_elems_5983], ==_arg_5980)
      let {cmps_5986 : [x_num_elems_5983]bool} =
        map(x_num_elems_5983,
            {x_flat_5984, y_flat_5985},
            \ {x_5987 : i32, y_5988 : i32}
              : {bool} ->
              let {binlam_res_5989 : bool} = eq_i32(x_5987, y_5988)
              in {binlam_res_5989})
      let {all_equal_5990 : bool} =
        redomap(x_num_elems_5983,
                {cmps_5986},
                {commutative \ {x_5991 : bool, y_5992 : bool}
                  : {bool} ->
                  let {binlam_res_5993 : bool} = logand(x_5991, y_5992)
                  in {binlam_res_5993},
                {true}},
                \ {x_5994 : bool}
                  : {bool} ->
                  {x_5994})
      in {all_equal_5990}
    } else {false}
    : {bool}
  let {&&_arg_5995 : bool} = eq_i64(n_5974, m_5975)
  let {main_res_5996 : bool} =
    if  &&_arg_5995
    then {arrays_equal_5982} else {
      let {x_5997 : bool} = false
      in {x_5997}
    }
    : {bool}
  in {main_res_5996}
}

entry("main",
      {xs: []i32,
       ys: []i32},
      {bool})
  entry_main (n_5998 : i64,
              m_5999 : i64,
              xs_6000 : [n_5998]i32,
              ys_6001 : [m_5999]i32)
  : {bool} = {
  let {entry_result_6002 : bool} =
    apply main_5952(n_5998, m_5999, xs_6000, ys_6001)
    : {bool}
  in {entry_result_6002}
}

Looks like this is the result of internalizing the arguments before calling eIf (they get internalized in withAutoMap), but I don't understand the workings of the size type system well enough to know why this leads to n being compared with n instead of m.

zfnmxt avatar Dec 01 '22 20:12 zfnmxt

It's comparing n and m, so I don't understand the problem.

athas avatar Dec 01 '22 21:12 athas

I'm an idiot and was just confused by dim_eq_5981, but I see the issue now---namely that the arguments, and hence the bodies of the conditional, are internalized before the conditional and so any statements generated during internalization are in the wrong scope in this example. It's fixed now, though, and everything in tests now passes. I'll make CI all green tomorrow.

zfnmxt avatar Dec 02 '22 01:12 zfnmxt

When I replace the prelude definition of map with

def map 'a [n] 'x (f: a -> x) (as: [n]a): *[n]x = f as

then type checking works, but internalisation (specifically defunctionalisation) goes into an infinite loop. How come?

athas avatar Dec 05 '22 13:12 athas

I've fixed the coercion issue; the only issue remaining I'm aware of when getting rid of the map intrinsic in the prelude is the infinite looping in the defunctionalization pass. It's caused by

https://github.com/diku-dk/futhark/blob/d700c152ccc240febfbf1e695a39b81106f0a5e4/src/Futhark/Internalise/Defunctionalise.hs#L1202

and

https://github.com/diku-dk/futhark/blob/d700c152ccc240febfbf1e695a39b81106f0a5e4/src/Futhark/Internalise/Defunctionalise.hs#L1267

zfnmxt avatar Dec 06 '22 01:12 zfnmxt

Getting closer to no-more-map-intrinsic salvation; only a few errors remaining for the tests in the root directory:

┌──────────┬────────┬────────┬───────────┐
│          │ passed │ failed │ remaining │
├──────────┼────────┼────────┼───────────┤
│ programs │ 480    │ 6      │ 0/486     │
├──────────┼────────┼────────┼───────────┤
│ runs     │ 511    │ 6      │ 0/517     │
└──────────┴────────┴────────┴───────────┘

Most of the failures are due to triggering the error here:

https://github.com/diku-dk/futhark/blob/0b789608136b7f809f256151a829c7b6e35190e1/src/Futhark/Internalise/AccurateSizes.hs#L43-L50

I'll investigate tomorrow.

zfnmxt avatar Dec 07 '22 05:12 zfnmxt

Yes! I was looking forward to this one. And the implementation is delightfully simple (assuming it works). Part of the purpose of the interpreter is to keep us honest wrt. simple semantics, and it looks like automap is indeed still simple.

athas avatar Jan 03 '23 16:01 athas

Yes! I was looking forward to this one. And the implementation is delightfully simple (assuming it works). Part of the purpose of the interpreter is to keep us honest wrt. simple semantics, and it looks like automap is indeed still simple.

I'm still kind of surprised that it actually (at least, apparently) works. I opted not to support short-circuiting for && and || if they're automapped---it would've just added complexity for something that's slow anyway.

zfnmxt avatar Jan 03 '23 19:01 zfnmxt

A function where automapping is ambiguous due to constraints imposed by size types:

def f [n][m][a] (A: [n][m][a]i32) : [n][a][m]i32 = transpose (transpose (transpose A))

This can either

transpose (transpose (map transpose A)

or

map transpose (transpose (transpose A))

In this case both would produce the same value, but I don't think that is the case in general. Such ambiguity should be flagged and produce an error.

athas avatar Feb 22 '23 09:02 athas

Magnus hints that this is relevant: https://www.cs.bu.edu/fac/snyder/publications/UnifChapter.pdf

The basic idea is that unification is possible (and unitary) when the underlying theory is associative, and shapes probably are. (Maybe - the Max constraint is a bit unclear, but maybe it's close enough.)

athas avatar Feb 22 '23 12:02 athas

In this case both would produce the same value, but I don't think that is the case in general. > Such ambiguity should be flagged and produce an error.

There are many such cases; although we could disambiguate this particular case by augmenting the minimal map rule with "maps are applied innermost first" rule or something. But that's probably too many rules for something that can already be very confusing.

Magnus hints that this is relevant: https://www.cs.bu.edu/fac/snyder/publications/UnifChapter.pdf

Reading this now; I think the current constraints now in the toy implementation are probably close to the "right thing". To review (sorry about the terrible infix names):

data Constraint
  = Eq Type Type
  | (:<) Type Type
  | SEq Shape Shape
  | (:$<) Shape Shape
  | Choice
      { try :: Constraints,
        backup :: Constraints
      }
  | HasShape Name Shape
  | Minimize Shape Shape Shape
  • Eq; For type equality.
  • SEq: For shape equality.
  • s1 :$< s2 : s1 must be a posfix of s2, i.e. s2 = s <> s1 where s is some shape.
  • t1 :< t2: Equivaent to shapeOf t1 :$< shapeOf t2 and Eq (baseType t1) (baseType t2) where baseType t returns the underlying non-array type of t.
  • Choice cs1 cs2: Try cs1, if it doesn't work, try cs2. Previously known as Max.
  • HasShape: Just a hack to assign a shape variable to type variables; never actually solved. Ignore.
  • Minimize s l u: Minimize the shape s, lower bounded by l and upper bounded by u.

Minimize is a real pain in the ass to solve in general and is used to enforce minimal maps. I guess it's time I become an expert in unification! I've also found it very difficult to detect failure vs. success. For example,

\A -> (A == 0) |> flatten

has type [m][n]S int -> []S int and you'll essentially get this with the system I detail below, but with unsolved constraints remaining.

Constraint generation for applications looks like this now:

constrain (App f e _) = do
  (f', t_f) <- constrain f
  (e', t_e) <- constrain e
  (a, b) <- case (distribute t_f) of
    (a :-> b) -> pure (a, b)
    (TVar fvar) -> do
      a <- newTVar
      b <- newTVar
      addConstraint $ Eq (TVar fvar) (a :-> b)
      pure (a, b)
    _ -> error ""
  let t_e' = distribute t_e
  e_frame <- newSVar
  b_frame <- newSVar
  s_e' <- shapeOf t_e'
  s_a <- shapeOf a
  addConstraint $
    Choice
      { try =
          [ a :< t_e',
            TArray b_frame a `Eq` t_e',
            Minimize b_frame mempty s_e'
          ],
        backup =
          [ t_e' :< a,
            TArray e_frame t_e' `Eq` a,
            b_frame `SEq` mempty,
            Minimize e_frame mempty s_a
          ]
      }
  pure (App f' e' (Id $ Type $ TArray b_frame b), TArray b_frame b)

The idea is:

f : a :-> b               (If `f` has type `S_f (a' :-> b')` the shape `S_f` is distributed ;
e : t_e                    so its type becomes `S_f a' :->S_f b' = a :-> b` ; this is legal because
a :< t_e                   type-level ""arrays"" of functions can only ever be values of the form `map f`.)
S_bframe a = t_e
-----------------------------------  map
f e : S_bframe b
f : a :-> b
e : t_e
t_e :< a
S_eframe t_e = a   (S_eframe appears here so that it can be minimized.)
-----------------------------------  replicate
f e : b

"Replicate" here is a bit of a misnomer, it really refers to maping function arguments, e.g.,

f : a -> a
ys : [n]a
ys |> f ---AUTOMAP--> ys |> map f 

But replication and mapping of function arguments is the same thing on the type level because we consider S (a :-> b) ~ S a :-> S b. But as a bonus the rule also lets us support replication-style broadcasting of non-function arguments (although I'm not sure we should support it since things like length 0 are gross).

In practice, the Choice constraint is used exclusively to choose between the map and replicate rules and all applications are typed via application of the map or reduce rules; both succeed exactly when the application is just standard function application.

zfnmxt avatar Feb 24 '23 02:02 zfnmxt

I don't really understand the implications of the second part. Maybe we should have a remote meeting about this.

athas avatar Feb 24 '23 07:02 athas