Merge pull request #1294 from axch/case-of-case

Case of case and other optimizations in the inliner

Author: axch

src/lib/Builder.hs

@@ -768,15 +768,15 @@ injectAltResult sumTys con (Abs b body) = liftBuilder do
 
 -- TODO: consider a version with nonempty list of alternatives where we figure
 -- out the result type from one of the alts rather than providing it explicitly
-buildCase :: (Emits n, ScopableBuilder r m)
-          => Atom r n -> Type r n
-          -> (forall l. (Emits l, DExt n l) => Int -> Atom r l -> m l (Atom r l))
-          -> m n (Atom r n)
-buildCase scrut resultTy indexedAltBody = do
+buildCase' :: (Emits n, ScopableBuilder r m)
+  => Atom r n -> Type r n
+  -> (forall l. (Emits l, DExt n l) => Int -> Atom r l -> m l (Atom r l))
+  -> m n (Expr r n)
+buildCase' scrut resultTy indexedAltBody = do
   case trySelectBranch scrut of
     Just (i, arg) -> do
       Distinct <- getDistinct
-      indexedAltBody i $ sink arg
+      Atom <$> indexedAltBody i (sink arg)
     Nothing -> do
       scrutTy <- getType scrut
       altBinderTys <- caseAltsBinderTys scrutTy
@@ -786,7 +786,13 @@ buildCase scrut resultTy indexedAltBody = do
           eff <- getEffects blk
           return $ blk `PairE` eff
         return (Abs b' body, ignoreHoistFailure $ hoist b' eff')
-      emitExpr $ Case scrut alts resultTy $ mconcat effs
+      return $ Case scrut alts resultTy $ mconcat effs
+
+buildCase :: (Emits n, ScopableBuilder r m)
+  => Atom r n -> Type r n
+  -> (forall l. (Emits l, DExt n l) => Int -> Atom r l -> m l (Atom r l))
+  -> m n (Atom r n)
+buildCase s r b = emitExprToAtom =<< buildCase' s r b
 
 buildEffLam
   :: ScopableBuilder r m

src/lib/CheapReduction.hs

@@ -15,7 +15,7 @@ module CheapReduction
   , unwrapLeadingNewtypesType, wrapNewtypesData, liftSimpAtom, liftSimpType
   , liftSimpFun, makeStructRepVal, NonAtomRenamer (..), Visitor (..), VisitGeneric (..)
   , visitAtomPartial, visitTypePartial, visitAtomDefault, visitTypeDefault, Visitor2
-  , visitBinders, visitPiDefault)
+  , visitBinders, visitPiDefault, visitAlt)
   where
 
 import Control.Applicative

src/lib/Inline.hs

@@ -16,6 +16,7 @@ import IRVariants
 import Name
 import Subst
 import Occurrence hiding (Var)
+import Optimize
 import Types.Core
 import Types.Primitives
 
@@ -90,7 +91,7 @@ inlineDeclsSubst = \case
       s <- getSubst
       extendSubst (b @> SubstVal (SuspEx expr s)) $ inlineDeclsSubst rest
     else do
-      expr' <- inlineExpr Stop expr
+      expr' <- inlineExpr Stop expr >>= (liftEnvReaderM . peepholeExpr)
       -- If the inliner starts moving effectful expressions, it may become
       -- necessary to query the effects of the new expression here.
       let presInfo = resolveWorkConservation ann expr'
@@ -248,6 +249,9 @@ data Context (from::E) (to::E) (o::S) where
   Stop :: Context e e o
   TabAppCtx :: [SAtom i] -> Subst InlineSubstVal i o
             -> Context SExpr e o -> Context SExpr e o
+  CaseCtx :: [SAlt i] -> SType i -> EffectRow SimpIR i
+          -> Subst InlineSubstVal i o
+          -> Context SExpr e o -> Context SExpr e o
   EmitToAtomCtx :: Context SAtom e o -> Context SExpr e o
   EmitToNameCtx :: Context SAtomName e o -> Context SAtom e o
 
@@ -271,6 +275,9 @@ inlineExpr ctx = \case
   TabApp tbl ixs -> do
     s <- getSubst
     inlineAtom (TabAppCtx ixs s ctx) tbl
+  Case scrut alts resultTy effs -> do
+    s <- getSubst
+    inlineAtom (CaseCtx alts resultTy effs s ctx) scrut
   expr -> visitGeneric expr >>= reconstruct ctx
 
 inlineAtom :: Emits o => Context SExpr e o -> SAtom i -> InlineM i o (e o)
@@ -340,12 +347,18 @@ instance Inlinable SBlock where
       effs' <- inline Stop effs  -- TODO Really?
       reconstruct ctx $ Block (BlockAnn ty' effs') decls' ans'
 
+inlineBlockEmits :: Emits o => Context SExpr e2 o -> SBlock i -> InlineM i o (e2 o)
+inlineBlockEmits ctx (Block _ decls ans) = do
+  inlineDecls decls $ inlineAtom ctx ans
+
 -- Still using InlineM because we may call back into inlining, and we wish to
 -- retain our output binding environment.
 reconstruct :: Emits o => Context e1 e2 o -> e1 o -> InlineM i o (e2 o)
 reconstruct ctx e = case ctx of
   Stop -> return e
   TabAppCtx ixs s ctx' -> withSubst s $ reconstructTabApp ctx' e ixs
+  CaseCtx alts resultTy effs s ctx' ->
+    withSubst s $ reconstructCase ctx' e alts resultTy effs
   EmitToAtomCtx ctx' -> emitExprToAtom e >>= reconstruct ctx'
   EmitToNameCtx ctx' -> emit (Atom e) >>= reconstruct ctx'
 {-# INLINE reconstruct #-}
@@ -404,5 +417,38 @@ reconstructTabApp ctx expr ixs =
       ixs' <- mapM (inline Stop) ixs
       reconstruct ctx $ TabApp array' ixs'
 
+reconstructCase :: Emits o
+  => Context SExpr e o -> SExpr o -> [SAlt i] -> SType i -> EffectRow SimpIR i
+  -> InlineM i o (e o)
+reconstructCase ctx scrutExpr alts resultTy effs =
+  case scrutExpr of
+    Case sscrut salts _ _ -> do
+      -- Perform case-of-case optimization
+      -- TODO Add join points to reduce code duplication (and repeated inlining)
+      -- of the arms of the outer case
+      resultTy' <- inline Stop resultTy
+      reconstruct ctx =<< (buildCase' sscrut resultTy' \i val -> do
+        ans <- applyAbs (sink $ salts !! i) (SubstVal val) >>= emitBlock
+        buildCase ans (sink resultTy') \j jval -> do
+          Abs b body <- return $ alts !! j
+          extendSubst (b @> (SubstVal $ DoneEx $ Atom jval)) do
+            inlineBlockEmits Stop body >>= emitExprToAtom)
+    _ -> do
+      -- Attempt case-of-known-constructor optimization
+      -- I can't use `buildCase` here because I want to propagate the incoming
+      -- context `ctx` into the selected alternative if the optimization fires,
+      -- but leave it around the whole reconstructed `Case` if it doesn't.
+      scrut <- emitExprToAtom scrutExpr
+      case trySelectBranch scrut of
+        Just (i, val) -> do
+          Abs b body <- return $ alts !! i
+          extendSubst (b @> (SubstVal $ DoneEx $ Atom val)) do
+            inlineBlockEmits ctx body
+        Nothing -> do
+          alts' <- mapM visitAlt alts
+          resultTy' <- inline Stop resultTy
+          effs' <- inline Stop effs
+          reconstruct ctx $ Case scrut alts' resultTy' effs'
+
 instance Inlinable (EffectRow SimpIR)
 instance Inlinable (EffectAndType SimpIR)

src/lib/Optimize.hs

@@ -7,7 +7,7 @@
 {-# LANGUAGE UndecidableInstances #-}
 
 module Optimize
-  ( optimize, peepholeOp
+  ( optimize, peepholeOp, peepholeExpr
   , hoistLoopInvariant, hoistLoopInvariantDest
   , dceTop, dceTopDest
   , foldCast ) where
@@ -43,7 +43,7 @@ optimize = dceTop     -- Clean up user code
 
 -- === Peephole optimizations ===
 
-peepholeOp :: PrimOp SimpIR o -> EnvReaderM o (Either (SAtom o) (PrimOp SimpIR o))
+peepholeOp :: PrimOp SimpIR o -> EnvReaderM o (SExpr o)
 peepholeOp op = case op of
   MiscOp (CastOp (BaseTy (Scalar sTy)) (Con (Lit l))) -> return $ case foldCast sTy l of
     Just l' -> lit l'
@@ -72,14 +72,14 @@ peepholeOp op = case op of
     return $ lit $ Word8Lit $ lv .|. rv
   BinOp BAnd (Con (Lit (Word8Lit lv))) (Con (Lit (Word8Lit rv))) ->
     return $ lit $ Word8Lit $ lv .&. rv
-  MiscOp (ToEnum ty (Con (Lit (Word8Lit tag)))) -> Left <$> case ty of
-    SumTy cases -> return $ SumVal cases (fromIntegral tag) UnitVal
+  MiscOp (ToEnum ty (Con (Lit (Word8Lit tag)))) -> case ty of
+    SumTy cases -> return $ Atom $ SumVal cases (fromIntegral tag) UnitVal
     _ -> error "Ill typed ToEnum?"
-  MiscOp (SumTag (SumVal _ tag _))                   -> return $ lit $ Word8Lit $ fromIntegral tag
+  MiscOp (SumTag (SumVal _ tag _)) -> return $ lit $ Word8Lit $ fromIntegral tag
   _ -> return noop
   where
-    noop = Right op
-    lit = Left . Con . Lit
+    noop = PrimOp op
+    lit = Atom . Con . Lit
 
     cmp :: Ord a => CmpOp -> a -> a -> Bool
     cmp = \case
@@ -188,9 +188,9 @@ foldCast sTy l = case sTy of
         compare (0 - countTrailingZeros (round @b @a a))
                 (0 - countTrailingZeros (round @b @a b))
 
-peepholeExpr :: SExpr o -> EnvReaderM o (Either (SAtom o) (SExpr o))
+peepholeExpr :: SExpr o -> EnvReaderM o (SExpr o)
 peepholeExpr expr = case expr of
-  PrimOp op -> fmap PrimOp <$> peepholeOp op
+  PrimOp op -> peepholeOp op
   TabApp (Var t) [IdxRepVal ord] ->
     lookupAtomName t <&> \case
       LetBound (DeclBinding ann _ (TabCon Nothing tabTy elems))
@@ -199,12 +199,12 @@ peepholeExpr expr = case expr of
         -- For example, it might be coming from an unsafe_from_ordinal that is
         -- under a case branch that would be dead for all invalid indices.
         if 0 <= ord && fromIntegral ord < length elems
-          then Left $ elems !! fromIntegral ord
-          else Right expr
-      _ -> Right expr
+          then Atom $ elems !! fromIntegral ord
+          else expr
+      _ -> expr
   -- TODO: Apply a function to literals when it has a cheap body?
   -- Think, partial evaluation of threefry.
-  _ -> return $ Right expr
+  _ -> return expr
   where isFinTabTy = \case
           TabPi (TabPiType (_:>(IxType _ (IxDictRawFin _))) _) -> True
           _ -> False
@@ -277,9 +277,8 @@ ulExpr expr = case expr of
   _ -> nothingSpecial
   where
     inc i = modify \(ULS n) -> ULS (n + i)
-    nothingSpecial = inc 1 >> (visitGeneric expr >>= liftEnvReaderM . peepholeExpr) >>= \case
-      Left x -> return x
-      Right e -> emitExpr e
+    nothingSpecial = inc 1 >> (visitGeneric expr >>= liftEnvReaderM . peepholeExpr)
+                     >>= emitExprToAtom
     unrollBlowupThreshold = 12
     withLocalAccounting m = do
       oldCost <- get

src/lib/Simplify.hs

@@ -758,9 +758,7 @@ simplifyGenericOp op = do
            (substM >=> getRepType)
            (simplifyAtom >=> toDataAtomIgnoreRecon)
            (error "shouldn't have lambda left")
-  result <- liftEnvReaderM (peepholeOp $ toPrimOp op') >>= \case
-    Left  a   -> return a
-    Right op'' -> emitOp op''
+  result <- liftEnvReaderM (peepholeOp $ toPrimOp op') >>= emitExprToAtom
   liftSimpAtom ty result
 {-# INLINE simplifyGenericOp #-}
 

src/lib/Types/Core.hs

@@ -425,6 +425,7 @@ type SAtom  = Atom SimpIR
 type SType  = Type SimpIR
 type SExpr  = Expr SimpIR
 type SBlock = Block SimpIR
+type SAlt   = Alt   SimpIR
 type SDecl  = Decl  SimpIR
 type SDecls = Decls SimpIR
 type SAtomName  = AtomName SimpIR

tests/inline-tests.dx

@@ -87,3 +87,44 @@ def id'(x:Nat) -> Nat = x
   sum (for i:(Fin 2) j:(..i). ordinal j)[ix]
 -- CHECK: 1
 -- CHECK-NOT: Compiler bug
+
+-- CHECK-LABEL: Inlining simplifies case-of-known-constructor
+"Inlining simplifies case-of-known-constructor"
+
+-- Inlining xs exposes a case-of-known-constructor opportunity here;
+-- the first inlining pass doesn't take it (yet) because it's
+-- conservative about inlining `i` into the body of `xs`, but the
+-- second pass does.
+%passes inline
+:pp
+  xs = for i:(Either (Fin 3) (Fin 4)).
+    case i of
+      Left k -> 1
+      Right k -> 2
+  for j:(Fin 3). xs[Left j]
+-- CHECK: === inline ===
+-- CHECK: for
+-- CHECK: case
+-- CHECK: === inline ===
+-- CHECK: for
+-- CHECK-NOT: case
+
+-- CHECK-LABEL: Inlining carries out the case-of-case optimization
+"Inlining carries out the case-of-case optimization"
+
+-- Before inlining there are two cases, but attempting to inline `x`
+-- reveals a case-of-case opprtunity, which in turn exposes
+-- case-of-known-constructor in each branch, leading to just one case
+-- in the end.
+%passes inline
+:pp
+  x = if id'(3) > 2
+        then Just 4
+        else Nothing
+  case x of
+    Just a -> a * a
+    Nothing -> 0
+-- CHECK: === inline ===
+-- CHECK: case
+-- CHECK-NOT: case
+-- CHECK: === inline ===