Taso rules

.gitignore

@@ -3,3 +3,5 @@ result.txt
 .stack-work/
 .envrc
 .direnv
+.vscode/
+dist-newstyle/

hie.yaml

@@ -9,6 +9,72 @@ cradle:
     - path: "./rules/debug/Paths_tensor_right.hs"
       component: "tensor-right:exe:rules-debug"
 
+    - path: "./rules/taso/concat/Main.hs"
+      component: "tensor-right:exe:rules-taso-concat"
+
+    - path: "./rules/taso/concat/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-concat"
+
+    - path: "./rules/taso/conv/Main.hs"
+      component: "tensor-right:exe:rules-taso-conv"
+
+    - path: "./rules/taso/conv/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-conv"
+
+    - path: "./rules/taso/enlarge/Main.hs"
+      component: "tensor-right:exe:rules-taso-enlarge"
+
+    - path: "./rules/taso/enlarge/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-enlarge"
+
+    - path: "./rules/taso/ewadd/Main.hs"
+      component: "tensor-right:exe:rules-taso-ewadd"
+
+    - path: "./rules/taso/ewadd/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-ewadd"
+
+    - path: "./rules/taso/ewmul/Main.hs"
+      component: "tensor-right:exe:rules-taso-ewmul"
+
+    - path: "./rules/taso/ewmul/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-ewmul"
+
+    - path: "./rules/taso/matmul2D/Main.hs"
+      component: "tensor-right:exe:rules-taso-matmul2D"
+
+    - path: "./rules/taso/matmul2D/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-matmul2D"
+
+    - path: "./rules/taso/matmul3D/Main.hs"
+      component: "tensor-right:exe:rules-taso-matmul3D"
+
+    - path: "./rules/taso/matmul3D/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-matmul3D"
+
+    - path: "./rules/taso/relu/Main.hs"
+      component: "tensor-right:exe:rules-taso-relu"
+
+    - path: "./rules/taso/relu/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-relu"
+
+    - path: "./rules/taso/smul/Main.hs"
+      component: "tensor-right:exe:rules-taso-smul"
+
+    - path: "./rules/taso/smul/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-smul"
+
+    - path: "./rules/taso/split/Main.hs"
+      component: "tensor-right:exe:rules-taso-split"
+
+    - path: "./rules/taso/split/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-split"
+
+    - path: "./rules/taso/transpose/Main.hs"
+      component: "tensor-right:exe:rules-taso-transpose"
+
+    - path: "./rules/taso/transpose/Paths_tensor_right.hs"
+      component: "tensor-right:exe:rules-taso-transpose"
+
     - path: "./rules/xla/add/Main.hs"
       component: "tensor-right:exe:rules-xla-add"
 

package.yaml

@@ -189,6 +189,73 @@ executables:
     dependencies: tensor-right
     ghc-options: *exe-ghc-options
     default-extensions: *exe-extensions
+  # TASO Executables
+  rules-taso-ewadd:
+    source-dirs: rules/taso/ewadd
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-ewmul:
+    source-dirs: rules/taso/ewmul
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-smul:
+    source-dirs: rules/taso/smul
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-relu:
+    source-dirs: rules/taso/relu
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-concat:
+    source-dirs: rules/taso/concat
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-transpose:
+    source-dirs: rules/taso/transpose
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-enlarge:
+    source-dirs: rules/taso/enlarge
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-matmul2D:
+    source-dirs: rules/taso/matmul2D
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-matmul3D:
+    source-dirs: rules/taso/matmul3D
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-split:
+    source-dirs: rules/taso/split
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
+  rules-taso-conv:
+    source-dirs: rules/taso/conv
+    main: Main.hs
+    dependencies: tensor-right
+    ghc-options: *exe-ghc-options
+    default-extensions: *exe-extensions
   # Other Executables
   rules-debug:
     source-dirs: rules/debug

plot/timing_plot.py

@@ -28,7 +28,7 @@ def first_num_of_tasks(self) -> int:
     def overall_time(self) -> float:
         return sum(x.time for x in self.results)
 
-
+# TODO: Handle ANSI color escape codes
 def parse_file(lines: Sequence[str]) -> list[Rule]:
     """
     The file looks like this:

rules/taso/concat/Main.hs

@@ -0,0 +1,95 @@
+{-# OPTIONS_GHC -Wno-missing-import-lists #-}
+
+import Grisette hiding ((-->))
+import TensorRight
+import TensorRight.Internal.DSL.TASO (concat, ewadd, ewmul, relu, smul)
+import Prelude hiding (concat)
+
+desugar :: forall a. NumRule a
+desugar _ = do
+  r <- newRClass "r"
+  [sa, sb] <- newMaps ["sa", "sb"] r
+  a <- newTensor @a "A" [r --> sa]
+  b <- newTensor @a "B" [r --> sb]
+  let d = ByRClass r
+  lhs <- concat d a b
+  rhs <- concatTensor a b d
+  rewrite "concat(d, A, B) ⇒ Concatenate((A, B), d)" lhs rhs
+
+smulAssociativity :: forall a. NumRule a
+smulAssociativity _ = do
+  let w = ("w" :: a)
+  r <- newRClass "r"
+  s <- newMap "s" r
+  x <- newTensor @a "x" [r --> s]
+  y <- newTensor @a "y" [r --> s]
+  let d = ByRClass r
+  lhs <- concat d (smul x w) (smul y w)
+  rhs <- smul (concat d x y) w
+  rewrite "concat(d, smul(x, w), smul(y, w)) ⇒ smul(concat(d, x, y), w)" lhs rhs
+
+ewaddAssociativity :: forall a. NumRule a
+ewaddAssociativity _ = do
+  r <- newRClass "r"
+  s <- newMap "s" r
+  x <- newTensor @a "x" [r --> s]
+  y <- newTensor @a "y" [r --> s]
+  z <- newTensor @a "z" [r --> s]
+  w <- newTensor @a "w" [r --> s]
+  let d = ByRClass r
+  lhs <- concat d (ewadd x y) (ewadd z w)
+  rhs <- ewadd (concat d x z) (concat d y w)
+  rewrite "concat(d, ewadd(x, y), ewadd(z, w)) ⇒ ewadd(concat(d, x, z), concat(d, y, w))" lhs rhs
+
+ewmulAssociativity :: forall a. NumRule a
+ewmulAssociativity _ = do
+  r <- newRClass "r"
+  s <- newMap "s" r
+  x <- newTensor @a "x" [r --> s]
+  y <- newTensor @a "y" [r --> s]
+  z <- newTensor @a "z" [r --> s]
+  w <- newTensor @a "w" [r --> s]
+  let d = ByRClass r
+  lhs <- concat d (ewmul x y) (ewmul z w)
+  rhs <- ewmul (concat d x z) (concat d y w)
+  rewrite "concat(d, ewmul(x, y), ewmul(z, w)) ⇒ ewmul(concat(d, x, z), concat(d, y, w))" lhs rhs
+
+reluAssociativity :: forall a. NumRule a
+reluAssociativity _ = do
+  r <- newRClass "r"
+  s <- newMap "s" r
+  x <- newTensor @a "x" [r --> s]
+  y <- newTensor @a "y" [r --> s]
+  let d = ByRClass r
+  lhs <- concat d (relu @a x) (relu @a y)
+  rhs <- relu @a $ concat d x y
+  rewrite "" lhs rhs
+
+geometry :: forall a. NumRule a
+geometry _ = do
+  [d0, d1, d2] <- newRClasses ["d0", "d1", "d2"]
+  d0S <- newMap "d0S" d0
+  d1S <- newMap "d1S" d1
+  d2S <- newMap "d2S" d2
+  x <- newTensor @a "x" [d0 --> d0S, d1 --> d1S, d2 --> d2S]
+  y <- newTensor @a "y" [d0 --> d0S, d1 --> d1S, d2 --> d2S]
+  z <- newTensor @a "z" [d0 --> d0S, d1 --> d1S, d2 --> d2S]
+  w <- newTensor @a "w" [d0 --> d0S, d1 --> d1S, d2 --> d2S]
+  lhs <- concat (ByRClass d0) (concat (ByRClass d1) x y) (concat (ByRClass d1) z w)
+  rhs <- concat (ByRClass d1) (concat (ByRClass d0) x z) (concat (ByRClass d0) y w)
+  rewrite "concat(d0, concat(d1, x, y), concat(d1, z, w)) ⇒ concat(d1, concat(d0, x, z), concat(0, y, w))" lhs rhs
+
+main :: IO ()
+main = do
+  printTitle "######################## desugarOneRole ########################"
+  verifyNumDSL desugar
+  printTitle "######################## smulAssociativity #####################"
+  verifyNumDSL smulAssociativity
+  printTitle "######################## ewaddAssociativity ####################"
+  verifyNumDSL ewaddAssociativity
+  printTitle "######################## ewmulAssociativity ####################"
+  verifyNumDSL ewmulAssociativity
+  printTitle "######################## reluAssociativity ####################"
+  verifyNumDSL reluAssociativity
+  printTitle "######################## geometry #############################"
+  verifyNumDSL geometry