Merge pull request #1277 from google-research/coalesce-dep-pairs

Coalesce dependent pair contexts in Imp representation.

Author: dougalm

lib/prelude.dx

@@ -2605,7 +2605,7 @@ def (**)(
     y: m=>n=>Float
     ) -> l=>n=>Float  given (l|Ix, m|Ix, n|Ix) =
   -- TODO(https://github.com/google-research/dex-lang/issues/1212) Replace with tiled_matmul.
-  naive_matmul(x, y)
+  tiled_matmul(x, y)
 
 def (**.)(mat: n=>m=>Float, v: m=>Float) -> (n=>Float) given (n|Ix, m|Ix) =
   for i. vdot(mat[i], v)

src/lib/Imp.hs

@@ -648,18 +648,30 @@ getElemTypeAndIdxStructure (LeafType ctxs baseTy) = case ctxs of
   Empty -> (UnboxedValue baseTy, Nothing)
   Nest b rest -> case b of
     TabCtx _ -> error "leading idxs should have been stripped off already"
-    DepPairCtx depBinder ->
-      case getIExprInterpretation (LeafType rest baseTy) of
+    DepPairCtx depBinder -> case splitLeadingDepPairs rest of
+      PairB depBinders rest' -> case getIExprInterpretation (LeafType rest' baseTy) of
         RawValue bt -> (UnboxedValue bt, Nothing)
         BufferPtr (BufferType ixs eltTy) -> do
-          let ixs' = case depBinder of
-                LeftB _      -> Nothing
-                RightB UnitB -> Just ixs
+          let ixs' = case allNothingBs (Nest depBinder depBinders) of
+                Just UnitB -> Just ixs
+                Nothing -> Nothing
           (BoxedBuffer eltTy, ixs')
     RefCtx -> (,Nothing) $ UnboxedValue $ hostPtrTy $ elemTypeToBaseType eltTy
       where BufferType _ eltTy = getRefBufferType (LeafType rest baseTy)
     where hostPtrTy ty = PtrType (CPU, ty)
 
+allNothingBs :: Nest (MaybeB b) n l -> Maybe (UnitB n l)
+allNothingBs Empty = Just UnitB
+allNothingBs (Nest (LeftB _) _) = Nothing
+allNothingBs (Nest (RightB UnitB) rest) = allNothingBs rest
+
+splitLeadingDepPairs :: TypeCtx SimpIR n l -> PairB (Nest (MaybeB SBinder)) (TypeCtx SimpIR) n l
+splitLeadingDepPairs = \case
+  Empty -> PairB Empty Empty
+  Nest (DepPairCtx b) rest -> case splitLeadingDepPairs rest of
+    PairB bs rest' -> PairB (Nest b bs) rest'
+  ctxs -> PairB Empty ctxs
+
 tryGetBoxIdxStructure :: LeafType n -> Maybe (IndexStructure SimpIR n)
 tryGetBoxIdxStructure leafTy = snd $ getElemTypeAndIdxStructure leafTy
 
@@ -727,14 +739,14 @@ valueToTree (RepVal tyTop valTop) = do
     RefTy _ t -> go (RNest ctx RefCtx) t val
     DepPairTy (DepPairType _ (b:>t1) (t2)) -> case val of
       Branch [v1, v2] -> do
-        case ctx of
-          REmpty -> do
+        case allDepPairCtxs (unRNest ctx) of
+          Just UnitB -> do
             tree1 <- rec t1 v1
             x <- repValAtom $ RepVal t1 v1
             t2' <- applySubst (b@>SubstVal x) t2
             tree2 <- go (RNest ctx (DepPairCtx NothingB )) t2' v2
             return $ Branch [tree1, tree2]
-          _ -> do
+          Nothing -> do
             tree1 <- rec t1 v1
             tree2 <- go (RNest ctx (DepPairCtx (JustB (b:>t1)))) t2 v2
             return $ Branch [tree1, tree2]
@@ -752,6 +764,11 @@ valueToTree (RepVal tyTop valTop) = do
     where rec = go ctx
 {-# INLINE valueToTree #-}
 
+allDepPairCtxs :: TypeCtx SimpIR n l -> Maybe (UnitB n l)
+allDepPairCtxs ctx = case splitLeadingDepPairs ctx of
+  PairB bs Empty -> allNothingBs bs
+  _ -> Nothing
+
 storeLeaf :: Emits n => LeafType n -> IExpr n -> IExpr n -> SubstImpM i n ()
 storeLeaf leafTy dest src = case getRefBufferType leafTy of
   BufferType Singleton (UnboxedValue _) -> store dest src

tests/eval-tests.dx

@@ -1118,6 +1118,20 @@ case frob of
 > 4.
 
 
+-- regression tests for #1212
+data Rectangle(a) = AsRectangle(n:Nat, m:Nat, elts:(Fin n => Fin m => a))
+data Brick(a) = AsBrick(n:Nat, m:Nat, l:Nat, elts:(Fin n => Fin m => Fin l => a))
 
+def mk_rect(n:Nat) -> Rectangle Nat =
+  AsRectangle _ _ $ for i:(Fin n) j:(Fin n). ordinal i
 
+def mk_brick(n:Nat) -> Brick Nat =
+  AsBrick _ _ _ $ for i:(Fin n) j:(Fin n) k:(Fin n). ordinal i
 
+rect = mk_rect(3)
+rect
+> (AsRectangle 3 3 [[0, 0, 0], [1, 1, 1], [2, 2, 2]])
+
+brick = mk_brick(3)
+brick
+> (AsBrick 3 3 3 [[[0, 0, 0], [0, 0, 0], [0, 0, 0]], [[1, 1, 1], [1, 1, 1], [1, 1, 1]], [[2, 2, 2], [2, 2, 2], [2, 2, 2]]])