map o map fusion bug in the presence of scatter

[coancea]

let main [n] (is : [n]i64) (ys_bar: *[n]f32) =
  let scatter_res_adj_gather =
      map(\ is_elem -> ys_bar[is_elem] ) is
  let zeros = replicate n 0.0f32
  let map_res_bar = scatter ys_bar is zeros
  let map_adjs_1 =
      map (\ lam_adj -> 5.0f32 * lam_adj ) scatter_res_adj_gather
  let map_adjs_2 =
      map (\ lam_adj -> 3.0f32 * lam_adj ) map_res_bar
  in (map_adjs_1, map_adjs_2)

generates:

Type error after pass 'Fuse SOACs':
In function main
Type error:
Variable ys_bar_4339 referenced after being consumed.

After some debugging, what happens is that the first map producing scatter_res_adj_gather is fused vertically with the second map producing map_adjs_1 (which is also fused horizontally with the third map producing map_adjs_2).

So the problem seems to be that the variables consumed by scatter are not preempted for further fusion, i.e., they are not treated in the same way in-place updates are. The following would fix it, but I did not yet meditate enough if it is a full solution (and need to reuse the code with the in-place case):

checkForUpdates res (Op (Futhark.Scatter _ _ _ written_info)) = do
  let updt_arrs = map (\(_,_,x)->x) written_info
  res' <-
    foldM addVarToInfusible res updt_arrs
  aliases <- mconcat <$> mapM varAliases updt_arrs
  let inspectKer k = k {inplace = aliases <> inplace k}
  return res' {kernels = M.map inspectKer $ kernels res'}

[coancea]

That is not enough: even with the fix above, on another example, it fuses an early map into a map following scatter, where the early map uses the array consumed by scatter as a free variable.

[athas]

Your fix was almost right, but due to the convoluted control flow in fusion, it wasn't ever calling checkForUpdates on Scatter SOACs.