Avoid duplicating first-order code in defunctionalisation.

This is done by eta-expanding uses of function-typed variables.  We
will still generate a wrapper function, but it will always be small,
while previously we would duplicate the entire first order code, which
can be enormous.  This will be particularly important once we stop
inlining so much.

Closes #1968.

CHANGELOG.md

@@ -55,6 +55,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 * Bug in alias checking for the core language type checker (#1949).
   Actually (finally) a proper fix of #803.
 
+* Defunctionalisation duplicates less code (#1968).
+
 ## [0.24.3]
 
 ### Fixed

src/Futhark/Internalise/Defunctionalise.hs

@@ -560,14 +560,15 @@ defuncExp (AppExp (Range e1 me incl loc) res) = do
 defuncExp e@(Var qn (Info t) loc) = do
   sv <- lookupVar (toStruct t) (qualLeaf qn)
   case sv of
-    -- If the variable refers to a dynamic function, we return its closure
-    -- representation (i.e., a record expression capturing the free variables
-    -- and a 'LambdaSV' static value) instead of the variable itself.
-    DynamicFun closure _ -> pure closure
+    -- If the variable refers to a dynamic function, we eta-expand it
+    -- so that we do not have to duplicate its definition.
+    DynamicFun {} -> do
+      (params, body, ret) <- etaExpand (RetType [] t) e
+      defuncFun [] params body ret mempty
     -- Intrinsic functions used as variables are eta-expanded, so we
     -- can get rid of them.
     IntrinsicSV -> do
-      (pats, body, tp) <- etaExpand (RetType [] (typeOf e)) e
+      (pats, body, tp) <- etaExpand (RetType [] t) e
       defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty
     HoleSV _ hole_loc ->
       pure (Hole (Info t) hole_loc, sv)
@@ -892,20 +893,24 @@ unRetType (RetType ext t) = do
 
 defuncApplyFunction :: Exp -> Int -> DefM (Exp, StaticVal)
 defuncApplyFunction e@(Var qn (Info t) loc) num_args = do
-  let (argtypes, _) = unfoldFunType t
+  let (argtypes, rettype) = unfoldFunType t
   sv <- lookupVar (toStruct t) (qualLeaf qn)
 
   case sv of
     DynamicFun _ _
       | fullyApplied sv num_args -> do
           -- We still need to update the types in case the dynamic
           -- function returns a higher-order term.
-          let (argtypes', rettype) = dynamicFunType sv argtypes
-          pure (Var qn (Info (foldFunType argtypes' $ RetType [] rettype)) loc, sv)
+          let (argtypes', rettype') = dynamicFunType sv argtypes
+          pure (Var qn (Info (foldFunType argtypes' $ RetType [] rettype')) loc, sv)
+      | all (orderZero . snd) argtypes,
+        orderZero rettype -> do
+          (params, body, ret) <- etaExpand (RetType [] t) e
+          defuncFun [] params body ret mempty
       | otherwise -> do
           fname <- newVName $ "dyn_" <> baseString (qualLeaf qn)
           let (pats, e0, sv') = liftDynFun (prettyString qn) sv num_args
-              (argtypes', rettype) = dynamicFunType sv' argtypes
+              (argtypes', rettype') = dynamicFunType sv' argtypes
               dims' = mempty
 
           -- Ensure that no parameter sizes are AnySize.  The internaliser
@@ -915,11 +920,11 @@ defuncApplyFunction e@(Var qn (Info t) loc) num_args = do
           let bound_sizes = S.fromList dims' <> globals
           pats' <- instAnySizes pats
 
-          liftValDec fname (RetType [] $ toStruct rettype) (dims' ++ unboundSizes bound_sizes pats') pats' e0
+          liftValDec fname (RetType [] $ toStruct rettype') (dims' ++ unboundSizes bound_sizes pats') pats' e0
           pure
             ( Var
                 (qualName fname)
-                (Info (foldFunType argtypes' $ RetType [] $ fromStruct rettype))
+                (Info (foldFunType argtypes' $ RetType [] $ fromStruct rettype'))
                 loc,
               sv'
             )

src/Futhark/Internalise/LiftLambdas.hs

@@ -97,15 +97,13 @@ liftFunction fname tparams params (RetType dims ret) funbody = do
       isSize (v, _) = v `M.member` unFV sizes_in_types
       (free_dims, free_nondims) = partition isSize free
 
-      free_params =
-        map (mkParam . second (`setUniqueness` Nonunique)) $
-          free_dims ++ free_nondims
+      free_ts = map (second (`setUniqueness` Nonunique)) $ free_dims ++ free_nondims
 
   addValBind $
     ValBind
       { valBindName = fname,
         valBindTypeParams = tparams,
-        valBindParams = free_params ++ params,
+        valBindParams = map mkParam free_ts ++ params,
         valBindRetDecl = Nothing,
         valBindRetType = Info (RetType dims ret),
         valBindBody = funbody,
@@ -117,7 +115,7 @@ liftFunction fname tparams params (RetType dims ret) funbody = do
 
   pure
     $ apply
-      (Var (qualName fname) (Info (augType $ free_dims ++ free_nondims)) mempty)
+      (Var (qualName fname) (Info (augType free_ts)) mempty)
     $ free_dims ++ free_nondims
   where
     orig_type = funType params $ RetType dims ret

tests/noinline/noinline6.fut

@@ -0,0 +1,8 @@
+-- ==
+-- input { [1,2,3] [4,5,6] } output { [5, 6, 7] [7, 9, 11] }
+-- structure gpu { SegMap/Apply 2 }
+
+#[noinline]
+def f y x = x + y + 2i32
+
+def main xs ys = (map (f 2) xs, map2 f xs ys)

tests/shapes/hof0.fut

@@ -1,7 +1,7 @@
 -- A dubious test - what we want to ensure is an absence of too many
 -- dynamic casts just after internalisation.
 -- ==
--- structure internalised { Assert 3 }
+-- structure internalised { Assert 2 }
 
 def f [k] 'a (dest: [k]a) (f: [k]a -> [k]a) : [k]a =
   f dest