linear folding

2 files changed, 52 insertions, 114 deletions

diff --git a/nn.in b/nn.in
index fb9a863..cfb5de9 100644
--- a/nn.in
+++ b/nn.in
@@ -1,79 +1,88 @@
 // Rules
-Symbolic(id) >< Add(out, b) => b ~ AddCheck(out, (*L)Symbolic(id));
+Linear(x, int q, int r) >< Add(out, b) => b ~ AddCheckLinear(out, x, q, r);
+
 Concrete(int k) >< Add(out, b)
 	| k == 0 => out ~ b
 	| _ => b ~ AddCheckConcrete(out, k);
-other >< Add(out, b) => out ~ Add((*L)other, b);
 
-Symbolic(id) >< AddCheck(out, a) => out ~ Add(a, (*L)Symbolic(id));
-Concrete(int j) >< AddCheck(out, a)
-	| j == 0 => out ~ a
-	| _ => out ~ Add(a, (*L)Concrete(j));
-other >< AddCheck(out, a) => out ~ Add(a, (*L)other);
+Linear(y, int s, int t) >< AddCheckLinear(out, x, int q, int r)
+	| q == 0 && r == 0 && s == 0 && t == 0 => out ~ Concrete(0), x ~ Eraser, y ~ Eraser
+	| s == 0 && t == 0 => Linear(x, q, r) ~ Materialize(out), y ~ Eraser
+	| q == 0 && r == 0 => (*L)Linear(y, s, t) ~ Materialize(out), x ~ Eraser
+	| _ => Linear(x, q, r) ~ Materialize(out_x), (*L)Linear(y, s, r) ~ Materialize(out_y), out ~ TermAdd(out_x, out_y);
+
+Concrete(int j) >< AddCheckLinear(out, x, int q, int r) => out ~ Linear(x, q, r + j);
+
+Linear(y, int s, int t) >< AddCheckConcrete(out, int k) => out ~ Linear(y, s, t + k);
 
-Symbolic(id) >< AddCheckConcrete(out, int k) => out ~ Add(Concrete(k), (*L) Symbolic(id));
 Concrete(int j) >< AddCheckConcrete(out, int k)
 	| j == 0 => out ~ Concrete(k)
 	| _ => out ~ Concrete(k + j);
-other >< AddCheckConcrete(out, int k) => out ~ Add(Concrete(k), (*L)other);
 
-Symbolic(id) >< Mul(out, b) => b ~ MulCheck(out, (*L)Symbolic(id));
+Linear(x, int q, int r) >< Mul(out, b) => b ~ MulCheck(out, x, q, r);
+
 Concrete(int k) >< Mul(out, b)
 	| k == 0 => b ~ Eraser, out ~ (*L)Concrete(0)
 	| k == 1 => out ~ b
 	| _ => b ~ MulCheckConcrete(out, k);
-other >< Mul(out, b) => out ~ Mul((*L)other, b);
 
-Symbolic(id) >< MulCheck(out, a) => out ~ Mul(a, (*L)Symbolic(id));
-Concrete(int j) >< MulCheck(out, a)
-  | j == 0 => a ~ Eraser, out ~ (*L)Concrete(0)
-	| j == 1 => out ~ a
-	| _ => out ~ Mul(a, (*L)Concrete(j));
-other >< MulCheck(out, a) => out ~ Mul(a, (*L)other);
+Linear(y, int s, int t) >< MulCheck(out, x, int q, int r)
+	| q == 0 && r == 0 && s == 0 && t == 0 => out ~ Concrete(0), x ~ Eraser, y ~ Eraser
+	| s == 0 && t == 0 => Linear(x, q, r) ~ Materialize(out), y ~ Eraser
+	| q == 0 && r == 0 => (*L)Linear(y, s, t) ~ Materialize(out), x ~ Eraser
+	| _ => Linear(x, q, r) ~ Materialize(out_x), (*L)Linear(y, s, r) ~ Materialize(out_y), out ~ TermMul(out_x, out_y);
+
+Concrete(int j) >< MulCheck(out, x, int q, int r) => out ~ Linear(x, q * j, r * j);
+
+Linear(y, int s, int t) >< MulCheckConcrete(out, int k) => out ~ Linear(y, s * k, t * k);
 
-Symbolic(id) >< MulCheckConcrete(out, int k) => out ~ Mul(Concrete(k), (*L)Symbolic(id));
 Concrete(int j) >< MulCheckConcrete(out, int k)
-  | j == 0 => out ~ Concrete(0)
+	| j == 0 => out ~ Concrete(0)
 	| j == 1 => out ~ Concrete(k)
 	| _ => out ~ Concrete(k * j);
-other >< MulCheckConcrete(out, int k) => out ~ Mul(Concrete(k), (*L)other);
 
-Symbolic(id) >< ReLU(out) => out ~ ReLU((*L)Symbolic(id));
+Linear(x, int q, int r) >< ReLU(out) => (*L)Linear(x, q, r) ~ Materialize(out_x), out ~ TermReLU(out_x);
 Concrete(int k) >< ReLU(out)
 	| k > 0 => out ~ (*L)Concrete(k)
 	| _ => out ~ Concrete(0);
-other >< ReLU(out) => out ~ ReLU((*L)other);
+
+Linear(x, int q, int r) >< Materialize(out)
+	| ? => out ~ Concrete(r), x ~ Eraser
+	| ? => out ~ Symbolic(x)
+	| ? => out ~ TermAdd(x, Concrete(r))
+	| ? => out ~ TermMul(Concrete(q), x)
+	| _ => out ~ TermAdd(TermMul(Concrete(q), x), Concrete(r));
 
 // Net
-Symbolic(X) ~ Add(out, b);
-Concrete(0) ~ Mul(b, Symbolic(Y));
-out;
+Linear(X, 1, 0) ~ Add(a, b);
+Concrete(0) ~ Mul(b, Linear(Y, 1, 0));
+a ~ Materialize(out);
+out; // Symbolic(X) 
+free out a b;
 
-free out;
-Symbolic(X) ~ Mul(out, Concrete(0));
-out;
+Linear(X, 1, 0) ~ Mul(a, Concrete(0));
+a ~ Materialize(out);
+out; // Concrete(0)
+free out a;
 
-free out;
-Symbolic(X) ~ Mul(out, b);
-Symbolic(Y) ~ Add(b, Concrete(1));
-out;
+Linear(X, 1, 0) ~ Mul(a, b);
+Linear(Y, 1, 0) ~ Add(b, Concrete(1));
+out; // Mul(Symbolic(X),Add(Symbolic(Y),Concrete(1)))
+free out a b;
 
-free out;
 Concrete(1) ~ Add(out, b);
 Concrete(2) ~ Add(b, c);
 Concrete(3) ~ Mul(c, Concrete(2));
-out;
+out; // Concrete(9)
+free out a b c;
 
-free out;
-Symbolic(X) ~ Mul(a, Symbolic(W1));
-Symbolic(Y) ~ Mul(b, Symbolic(W2));
-b ~ Add(c, Symbolic(B));
+Linear(X, 1, 0) ~ Mul(a, Linear(W1, 1, 0));
+Linear(Y, 1, 0) ~ Mul(b, Linear(W2, 1, 0));
+b ~ Add(c, Linear(B, 1, 0));
 a ~ Add(d, c);
 d ~ ReLU(out);
-out;
+out; // ReLU(Add(Mul(Symbolic(X),Symbolic(W1)),Add(Mul(Symbolic(Y),Symbolic(W2)),Symbolic(B))))
+free out a b c d;
 
-free out;
-Symbolic(X) ~ Dup(a, b);
-a ~ Add(c, Concrete(1));
-b ~ Mul(out, c);
+Linear(X, 2, 10) ~ Materialize(out);
 out;
diff --git a/nn.in~ b/nn.in~
deleted file mode 100644
index dd89d03..0000000
--- a/nn.in~
+++ /dev/null
@@ -1,71 +0,0 @@
-// Agents
-
-Any >< Optimize(out) =>
-	out ~ Any;
-
-// Rules
-
-// Optimizer interacts with Const
-Const(k) >< Optimize(out) =>
-	out ~ Const(k);
-
-
-// Optimizer interacts with Add
-Add(a, b) >< Optimize(out) =>
-	a ~ Optimize(a_opt),
-	a_opt ~ Add_CheckLeft(out, b);
-
-// CHECK LEFT
-Const(int k) >< Add_CheckLeft(out, b)
-	| k == 0 => b ~ Optimize(b_opt), out ~ b_opt
-	| _ => b ~ Optimize(b_opt), b_opt ~ Add_CheckRight_WithConst(out, k);
-a_opt >< Add_CheckLeft(out, b) =>
-	b ~ Optimize(b_opt),
-	b_opt ~ Add_CheckRight_WithAny(out, (*L)a_opt);
-
-// CHECK RIGHT
-Const(int j) >< Add_CheckRight_WithConst(out, int k) =>
-	out ~ Const(k + j);
-b_opt >< Add_CheckRight_WithConst(out, int k) =>
-	out ~ Add(Const(k), (*L)b_opt);
-Const(int j) >< Add_CheckRight_WithAny(out, a_opt)
-	| j == 0 => out ~ a_opt
-	| _ => out ~ Add(a_opt, Const(j));
-b_opt >< Add_CheckRight_WithAny(out, a_opt) =>
-	out ~ Add(a_opt, (*L)b_opt);
-
-
-// Optimizer interacts with Mul
-Mul(a, b) >< Optimize(out) =>
-	a ~ Optimize(a_opt),
-	a_opt ~ Mul_CheckLeft(out, b);
-
-// Check LEFT
-Const(int k) >< Mul_CheckLeft(out, b)
-	| k == 0 => b ~ Eraser, out ~ Const(0)
-	| k == 1 => b ~ Optimize(b_opt), out ~ b_opt
-	| _ => b ~ Optimize(b_opt), b_opt ~ Mul_CheckRight_WithConst(out, k);
-a_opt >< Mul_CheckLeft(out, b) =>
-	b ~ Optimize(b_opt),
-	b_opt ~ Mul_CheckRight_WithAny(out, (*L)a_opt);
-
-// Check RIGHT
-Const(int j) >< Mul_CheckRight_WithConst(out, int k) =>
-	out ~ Const(k * j);
-b_opt >< Mul_CheckRight_WithConst(out, int k) =>
-	out ~ Mul(Const(k), (*L)b_opt);
-Const(int j) >< Mul_CheckRight_WithAny(out, a_opt)
-	| j == 0 => Eraser ~ a_opt, out ~ Const(0)
-	| j == 1 => out ~ a_opt
-	| _ => out ~ Mul(a_opt, Const(j));
-b_opt >< Mul_CheckRight_WithAny(out, a_opt) =>
-	out ~ Mul(a_opt, (*L)b_opt);
-
-
-// Net
-Mul(Const(1), Add(Any, Mul(Const(0), Add(Any, Any)))) ~ Optimize(out);
-out;
-
-free out;
-Mul(Const(0), Add(Any, Any)) ~ Optimize(out);
-out;