smart matrix multiplication (not enabled) and common factor detection

Author: davmre

matrizer.cabal

@@ -28,7 +28,7 @@ executable matrizer
 
 
 Library
-  Build-Depends:     base >=4.5 && < 5, containers, mtl >=2.1 && <2.2, parsec >=3.1.3 && <3.2, multimap >=1.2, transformers >=0.3 && <0.4, heap >= 1.0.2, filepath, directory
+  Build-Depends:     base >=4.5 && < 5, containers, mtl >=2.1 && <2.2, parsec >=3.1.3 && <3.2, multimap >=1.2, transformers >=0.3 && <0.4, heap >= 1.0.2, filepath, directory, array
   hs-source-dirs:    src
   Exposed-modules:   Matrizer.MTypes, Matrizer.Parsing, Matrizer.Optimization, Matrizer.Analysis, Matrizer.CodeGen, Matrizer.Util, Matrizer.Search, Matrizer.Preprocess, Matrizer.Derivatives, Matrizer.RewriteRules, Matrizer.Equivalence
   default-language:  Haskell2010

src/Matrizer/Derivatives.hs

@@ -12,6 +12,8 @@ import Matrizer.Optimization
 import Matrizer.Analysis
 import Matrizer.Search
 
+import Debug.Trace
+
 llexpr1 = (Branch3 MTernaryProduct (Branch1 MTranspose (Leaf "w")) (Branch2 MProduct (Branch1 MTranspose (Leaf "X")) (Leaf "X") ) (Leaf "w"))
 llexpr2 = (Branch2 MScalarProduct (LiteralScalar 2) (Branch2 MProduct (Branch1 MTranspose (Leaf "w")) (Branch2 MProduct (Branch1 MTranspose (Leaf "X")) (Leaf "y"))))
 llexpr3 = (Branch2 MProduct (Branch1 MTranspose (Leaf "y")) (Leaf "y"))
@@ -167,18 +169,15 @@ beamSearch2 fn iters beamSize nRewrites tbl beam =
 
 
 llSymbols2 :: SymbolTable
-llSymbols2 = Map.fromList [("X", Matrix 100 100 []), ("B", Matrix 100 100 []),("A", Matrix 100 100 []),("S", Matrix 100 100 []),("y", Matrix 100 1 []), ("C", Matrix 100 100 [])]
+llSymbols2 = Map.fromList [("X", Matrix 100 60 []), ("B", Matrix 100 100 []),("A", Matrix 100 100 []),("S", Matrix 60 60 []),("y", Matrix 100 1 []), ("C", Matrix 100 100 [])]
 
 pp = (Branch3 MTernaryProduct (Leaf "X") (Leaf "S") (Branch1 MTranspose (Leaf "X")))
 ll17 = (Branch1 MTrace (Branch3 MTernaryProduct (Leaf "A") (Branch1 MInverse pp) (Leaf "B")))
 dl17 = reduceDifferential "X" ll17
 
 
-
-dl17_1 = Branch1 MTrace (Branch2 MProduct (Leaf "A") (Branch2 MProduct (Branch3 MTernaryProduct (Leaf "C") (Branch2 MProduct (Branch1 MDifferential (Leaf "X")) (Branch2 MProduct (Leaf "S") (Branch1 MTranspose (Leaf "X")))) (Leaf "C")) (Leaf "B")))
-dl17_2 = Branch1 MTrace (Branch2 MProduct (Leaf "A") (Branch2 MProduct (Branch3 MTernaryProduct (Leaf "C")  (Branch2 MProduct (Leaf "X") (Branch2 MProduct (Leaf "S") (Branch1 MTranspose (Branch1 MDifferential (Leaf "X"))))) (Leaf "C")) (Leaf "B")))
-
-
+llS = Map.fromList [("X", Matrix 100 2 []), ("w", Matrix 2 1 []),("d", Matrix 2 1 []) ]
+r = (Branch2 MSum (Branch2 MProduct (Branch2 MProduct (Branch2 MProduct (Branch1 MTranspose (Leaf "d")) (Branch1 MTranspose (Leaf "X"))) (Leaf "X")) (Leaf "w")) (Branch2 MProduct (Branch2 MProduct (Branch2 MProduct (Branch1 MTranspose (Leaf "w")) (Branch1 MTranspose (Leaf "X"))) (Leaf "X")) (Leaf "d")))
 
 decompose :: Expr -> [(Float, Expr)]
 decompose (Branch2 MScalarProduct (LiteralScalar c) a) = [(c1*c, d) | (c1, d) <- decompose a]
@@ -214,7 +213,7 @@ differentiate tbl (Branch2 MDiff a b) c = do da <- differentiate tbl a c
 differentiate tbl (Branch2 MScalarProduct (LiteralScalar s) a) c = 
               do da <- differentiate tbl a c
                  return $ (Branch2 MScalarProduct (LiteralScalar s) da)
-differentiate tbl expr c =  differentiateBySearch tbl expr c
+differentiate tbl expr c =   differentiateBySearch tbl expr c
 
 
 differentiateBySearch tbl expr c = 
@@ -227,6 +226,8 @@ differentiateBySearch tbl expr c =
 derivFromAstar :: SymbolTable -> Expr -> Maybe Expr
 derivFromAstar tbl expr = do r <- runAstar tbl expr
                              extractDeriv tbl (nvalue r)
+
+
 
 reduceWithTrace tbl expr c = let d = reduceDifferential c expr 
                                  Right (Matrix d1 d2 dprops) = treeMatrix d tbl in

src/Matrizer/MTypes.hs

@@ -18,9 +18,9 @@ data Expr = Leaf VarName
            | Branch2 BinOp Expr Expr
            | Branch3 TernOp Expr Expr Expr
            | Let VarName Expr Bool Expr -- bool flag specifies whether this intermediate variable can be optimized out
-           deriving (Eq, Ord)
+           deriving (Eq, Ord, Show)
 
-data TernOp = MTernaryProduct deriving (Eq, Ord)
+data TernOp = MTernaryProduct deriving (Eq, Ord, Show)
 data BinOp = MProduct
            | MSum
            | MDiff
@@ -30,7 +30,7 @@ data BinOp = MProduct
            | MScalarProduct
            | MHadamardProduct
            | MColProduct
-           deriving (Eq, Ord, Enum)
+           deriving (Eq, Ord, Enum, Show)
 
 data UnOp = MInverse
           | MTranspose
@@ -44,60 +44,58 @@ data UnOp = MInverse
           | MDiagMV -- extract a matrix diagonal as a vector
           | MEntrySum
           | MElementWise ScalarOp
-          deriving (Eq, Ord)
+          deriving (Eq, Ord, Show)
 
 data ScalarOp = MLog
               | MExp -- TODO: support matrix exponentials
               | MReciprocal
-              deriving (Eq, Ord, Enum)
+              deriving (Eq, Ord, Enum, Show)
 
 -- AST pretty printing
+pprint_ternop MTernaryProduct = "***"
 
-instance Show TernOp where
-    show _ = "***"
-
-instance Show BinOp where
-    show MProduct = "mmul"
-    show MScalarProduct = "smul"
-    show MHadamardProduct = "hmul"
-    show MColProduct = "cmul" -- don't really expect people to use this input syntax 
+pprint_binop MProduct = "mmul"
+pprint_binop MScalarProduct = "smul"
+pprint_binop MHadamardProduct = "hmul"
+pprint_binop MColProduct = "cmul" -- don't really expect people to use this input syntax 
                             -- except for internal test cases
-    show MSum = "add"
-    show MDiff = "sub"
-    show MLinSolve = "solve"
-    show MTriSolve = "triSolve"
-    show MCholSolve = "cholSolve"
-
-instance Show UnOp where
-    show MInverse = "inv"
-    show MTranspose = "transpose"
-    show MChol = "chol"
-    show MTrace = "tr"
-    show (MDeriv v) = "deriv_" ++ v
-    show (MUnresolvedDeriv v) = "unresolved_deriv_" ++ v
-    show MDifferential = "differential"
-    show MDet = "det"
-    show MDiagVM = "toDiag"
-    show MDiagMV = "diag"
-    show MEntrySum = "sum"
-    show (MElementWise sop) = show sop
-
-instance Show ScalarOp where
-    show MLog = "log"
-    show MExp = "exp"
-    show MReciprocal = "recip"
-
-instance Show Expr where
-    show (Leaf a) = a
-    show (IdentityLeaf _) = "I"
-    show (ZeroLeaf _ _) = "0"
-    show (LiteralScalar x) = show x
-    show (Branch1 op c) = "(" ++ show op ++ " " ++ show c ++ ")"
-    show (Branch2 op a b) = "(" ++ show op ++ " " ++ show a ++ " "
-         ++ show b ++ ")"
-    show (Branch3 op a b c) = "(" ++ show op ++ " " ++ show a ++ " "
-         ++ show b ++ " " ++ show c ++ ")"
-    show (Let v a tmp b) = "(let (" ++ v ++ " := " ++ show a ++ (if tmp then " #temporary ) "  else ") ") ++ "\n" ++ show b ++ ")"
+pprint_binop MSum = "add"
+pprint_binop MDiff = "sub"
+pprint_binop MLinSolve = "solve"
+pprint_binop MTriSolve = "triSolve"
+pprint_binop MCholSolve = "cholSolve"
+
+
+pprint_unop MInverse = "inv"
+pprint_unop MTranspose = "transpose"
+pprint_unop MChol = "chol"
+pprint_unop MTrace = "tr"
+pprint_unop (MDeriv v) = "deriv_" ++ v
+pprint_unop (MUnresolvedDeriv v) = "unresolved_deriv_" ++ v
+pprint_unop MDifferential = "differential"
+pprint_unop MDet = "det"
+pprint_unop MDiagVM = "toDiag"
+pprint_unop MDiagMV = "diag"
+pprint_unop MEntrySum = "sum"
+pprint_unop (MElementWise sop) = pprint_scalarop sop
+
+pprint_scalarop MLog = "log"
+pprint_scalarop MExp = "exp"
+pprint_scalarop MReciprocal = "recip"
+
+
+pprint (Leaf a) = a
+pprint (IdentityLeaf _) = "I"
+pprint (ZeroLeaf _ _) = "0"
+pprint (LiteralScalar x) = show x
+pprint (Branch1 op c) = "(" ++ pprint_unop op ++ " " ++ pprint c ++ ")"
+pprint (Branch2 op a b) = "(" ++ pprint_binop op ++ " " ++ pprint a ++ " "
+         ++ pprint b ++ ")"
+pprint (Branch3 op a b c) = "(" ++ pprint_ternop op ++ " " ++ pprint a ++ " "
+         ++ pprint b ++ " " ++ pprint c ++ ")"
+pprint (Let v a tmp b) = "(let (" ++ v ++ " := " ++ pprint a ++ (if tmp then " #temporary ) "  else ") ") ++ "\n" ++ pprint b ++ ")"
+
+
 
 ------------------------------------------------------------------------
 -- Symbol Table Definition

src/Matrizer/Optimization.hs

@@ -11,6 +11,7 @@ import Control.Monad.Error
 import Matrizer.MTypes
 import Matrizer.Analysis
 import Matrizer.RewriteRules
+import Debug.Trace
 
 ----------------------------------
 
@@ -343,7 +344,7 @@ optimizerTraversal fn tbl rules z@( n@(Branch1 _ _), _) =
 -- the net change in FLOPs
 optimizeAtNode :: ScoreFn -> SymbolTable -> [Rule] -> Expr -> ThrowsError [(Expr, Int)]
 optimizeAtNode fn tbl rules t = let opts = mapMaybeFunc t [f tbl | f <- rules ] in
-                       scoreOptimizations fn tbl t opts
+                                 scoreOptimizations fn tbl t opts
 
 scoreOptimizations :: ScoreFn -> SymbolTable -> Expr -> [Expr] -> ThrowsError [(Expr, Int)]
 scoreOptimizations fn tbl t opts = mapM (scoreOpt t) opts where
@@ -373,3 +374,4 @@ mapMaybeFunc x (f:fs) =
 
 recognizeVar :: VarName -> Expr -> SymbolTable -> Expr -> Maybe Expr
 recognizeVar var rhs tbl tree = if rhs==tree then Just (Leaf var) else Nothing
+

src/Matrizer/RewriteRules.hs

@@ -3,13 +3,23 @@ module Matrizer.RewriteRules (
  Rules,
  optimizationRules,
  optimizationRulesFull,
+ rotateTraceLeft,
+ rotateTraceRight,
+ productsToList,
+ smartCommonFactor,
+ optimalProduct,
+ triplePartition,
+ commonPrefix
 ) where 
 
 import Data.Maybe
+import Data.Array
 
 import Matrizer.MTypes
 import Matrizer.Analysis
 
+import Debug.Trace
+
 ------------------------------------------------------------------
 -- List of optimizations
 --
@@ -39,12 +49,12 @@ type Rules = [Rule]
 
 optimizationRules :: Rules
 optimizationRules = inverseRules ++ transposeRules ++ binopSumRules ++
-    binopProductRules ++ ternProductRules ++ letExpRules ++ traceRules ++ detRules ++ diagRules ++ entrySumRules ++ hadamardProductRules ++ elementWiseRules
+    binopProductRules ++ ternProductRules ++ letExpRules ++ traceRules ++ detRules ++ diagRules ++ entrySumRules ++ hadamardProductRules ++ elementWiseRules  ++ [smartCommonFactor]
 
 -- moves that are valid and sometimes necessary, but generate many
 -- matches and can slow down inference. 
 expensiveMoves :: Rules
-expensiveMoves = [introduceTranspose]
+expensiveMoves = [introduceTranspose, productRtL, productLtR]
 
 optimizationRulesFull = optimizationRules ++ expensiveMoves
 
@@ -107,6 +117,8 @@ traceRules = [dissolveTrace
               , linearTrace
               , identityOps
               , traceProduct
+              , rotateTraceLeft
+              , rotateTraceRight
               , traceDiag]
 
 detRules :: Rules
@@ -167,6 +179,135 @@ assocMult _ (Branch2 MScalarProduct (Branch2 MScalarProduct l c) r) = Just (Bran
 assocMult _ (Branch2 MScalarProduct l (Branch2 MScalarProduct c r)) = Just (Branch2 MScalarProduct (Branch2 MScalarProduct l c) r)
 assocMult _ _ = Nothing
 
+
+-----------------------------------------------------------------------
+-- 'smart' product rules that collapse an entire chain of products (a*b*c*d*e*....) into a list,
+-- then re-expand that list
+-- TODO: write code for the optimal re-expansion
+
+productsToList :: Expr -> [Expr]
+productsToList (Branch2 MProduct a b) = (productsToList a) ++ (productsToList b)
+productsToList (Branch3 MTernaryProduct a b c) = (productsToList a) ++ (productsToList b) ++ (productsToList c)
+productsToList a = [a]
+
+expandProductLtR :: [Expr] -> Expr
+expandProductLtR (x:[]) = x
+expandProductLtR (x:xs) = (Branch2 MProduct x (expandProductLtR xs))
+
+-- takes a reversed list
+expandProductRtL :: [Expr] -> Expr
+expandProductRtL (x:[]) = x
+expandProductRtL (x:xs) = (Branch2 MProduct (expandProductRtL xs) x)
+
+productLtR :: Rule 
+productLtR _ z@(Branch2 MProduct (Branch2 MProduct _ _) _) = Just $ expandProductLtR $ productsToList z
+productLtR _ z@(Branch2 MProduct _ (Branch2 MProduct _ _)) = Just $ expandProductLtR $ productsToList z
+productLtR _ z@(Branch3 MTernaryProduct _ _ _) = Just $ expandProductLtR $ productsToList z
+productLtR _ _ = Nothing
+
+productRtL :: Rule 
+productRtL _ z@(Branch2 MProduct (Branch2 MProduct _ _) _) = Just $ expandProductRtL $ reverse $ productsToList z
+productRtL _ z@(Branch2 MProduct _ (Branch2 MProduct _ _)) = Just $ expandProductRtL $ reverse $ productsToList z
+productRtL _ z@(Branch3 MTernaryProduct _ _ _) = Just $ expandProductRtL $ reverse $ productsToList z
+productRtL _ _ = Nothing
+
+
+-- flopCost, split, rows, cols
+optimalProductArray :: [(Int, Int)] -> Array (Int, Int) (Int, Int, Int, Int)
+optimalProductArray  sizeList = r
+ where 
+       n = length sizeList
+       (rows, cols) = unzip sizeList
+       r = listArray ((0,0),(n,n))  (map oP (range ((0,0), (n,n))))
+       oP (i, j) = if i==j then (0, i, 0, 0)
+                   else if (j==i+1) then (0, i, rows!!i, cols!!i)
+                   else let costs = [(prodCost (r!(i,k)) (r!(k,j)), k) | k <- [(i+1)..(j-1)]] 
+                            (minCost, k) = minimum costs in
+                        (minCost, k, rows!!i, cols!!(j-1))
+       prodCost (f1, k1, r1, c1) (f2, k2, r2, c2) = f1 + f2 + (r1*c2 * (2*c1-1))
+
+optimalProduct tbl exprList = 
+    case length exprList of
+    1 -> head exprList
+    2 -> (Branch2 MProduct (exprList!!0) (exprList!!1) )
+    _ -> let sizeList = map exprSize exprList
+             r = optimalProductArray sizeList in
+             assembleProduct exprList r 0 (length exprList)
+ where exprSize e = let (Right (Matrix a b _)) = treeMatrix e tbl in
+                        (a, b)
+       assembleProduct exps r i j = if j==i+1 then exps!!i
+                                    else let (flops, k, a, b) =  r!(i,j) in 
+                                         (Branch2 MProduct (assembleProduct exps r i k) 
+                                                           (assembleProduct exps r k j))
+
+
+
+
+
+smartCommonFactor tbl (Branch2 MSum a@(Branch2 MProduct _ _) b@(Branch2 MProduct _ _)) = scfHelperOuter tbl a b
+smartCommonFactor tbl (Branch2 MSum a@(Branch3 MTernaryProduct _ _ _) b@(Branch2 MProduct _ _)) = scfHelperOuter tbl a b
+smartCommonFactor tbl (Branch2 MSum a@(Branch2 MProduct _ _) b@(Branch3 MTernaryProduct _ _ _)) = scfHelperOuter tbl a b
+smartCommonFactor tbl (Branch2 MSum a@(Branch3 MTernaryProduct _ _ _) b@(Branch3 MTernaryProduct _ _ _)) = scfHelperOuter tbl a b
+smartCommonFactor _ _ = Nothing
+
+
+scfHelperOuter tbl a b = 
+          let list1 = productsToList a 
+              list2 = productsToList b in
+           if list1==list2 then Just (Branch2 MScalarProduct (LiteralScalar 2.0) a)
+           else let (ncommon, clist) = (scfHelper tbl list1 list2) in
+                    if ncommon == 0 then Nothing
+                    else Just (optimalProduct tbl clist) 
+
+-- divide two lists into a common prefix, common suffix, and 'cores' that are different.
+-- for example, [1,4,8,2,3] and [1,4,0,3] becomes prefix=[1,4], cores=([8,2], [0]), suffix=[3] 
+triplePartition a b =  let (start, rA, rB) = commonPrefix a b
+                           (rend, rcoreA, rcoreB) = commonPrefix (reverse rA) (reverse rB) 
+                           (end, coreA, coreB) = (reverse rend, reverse rcoreA, reverse rcoreB) in
+                           (start, (coreA, coreB), end)
+commonPrefix (x:xs) (y:ys) = if x==y 
+                                then let (p1, r1, r2) = commonPrefix xs ys in
+                                     ((x:p1), r1, r2)
+                                else ([], (x:xs), (y:ys))
+commonPrefix a []  = ([], a, [])                                    
+commonPrefix [] b  = ([], [], b)
+
+-- we assume the lists are unequal, so at most one of the cores can be empty
+scfHelper tbl list1 list2 =
+                    let (start, (coreA, coreB), end) = triplePartition list1 list2
+                        (cA, cB) = (optProdOrEmpty tbl coreA coreB, optProdOrEmpty tbl coreB coreA) in
+                        (((length start) + (length end)), start ++ [Branch2 MSum  cA  cB] ++ end)
+   where                        
+       optProdOrEmpty tbl [] (b:bs) = let Right (Matrix r1 r2 []) = treeMatrix b tbl in
+                                          (IdentityLeaf r1)
+       optProdOrEmpty tbl a _ = optimalProduct tbl a
+          
+
+rotateTraceLeft :: Rule
+rotateTraceLeft tbl (Branch1 MTrace z@(Branch2 MProduct _ _)) = rtlHelper tbl z
+rotateTraceLeft tbl (Branch1 MTrace z@(Branch3 MTernaryProduct _ _ _)) = rtlHelper tbl z
+rotateTraceLeft _ _ = Nothing
+rtlHelper tbl z = 
+                let x:xs = productsToList z 
+                    rotated = xs ++ [x]
+                    candidate = expandProductRtL $ reverse rotated in
+                    case treeMatrix  candidate tbl of
+                    Right m -> Just (Branch1 MTrace candidate)
+                    Left err -> Nothing
+
+rotateTraceRight :: Rule
+rotateTraceRight tbl (Branch1 MTrace z@(Branch2 MProduct _ _)) = rtrHelper tbl z
+rotateTraceRight tbl (Branch1 MTrace z@(Branch3 MTernaryProduct _ _ _)) = rtrHelper tbl z
+rotateTraceRight _ _ = Nothing
+rtrHelper tbl z = 
+                let x:xs = reverse $ productsToList z 
+                    candidate = expandProductRtL $ (xs ++ [x]) in
+                    case treeMatrix  candidate tbl of
+                    Right m -> Just (Branch1 MTrace candidate)
+                    Left err -> Nothing
+                    
+-------------------------------------------------
+
 -- (AC + BC) -> (A+B)C
 commonFactorRight :: Rule
 commonFactorRight _ (Branch2 MSum (Branch2 MProduct l1 l2) (Branch2 MProduct r1 r2)) =