Bytes chip optimization and new operations (#421)

* `u8_chain_rotr7` and `u8_chain_rotr4` added

* removed `bit_decomposition` from `blake3_g_function`

* `u8_add` now outputs a single column

the carry can be derived from the expression `(x + y - z)/256`

* `u8_sub` now outputs single column

* cross test

Author: gabriel-barrett

Ix/Aiur/Compiler/Check.lean

@@ -739,6 +739,14 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
     let a' ← checkNoEscape a .field
     let b' ← checkNoEscape b .field
     pure (Typed.Term.u8Mul (.tuple #[.field, .field]) false a' b')
+  | .u8ChainRotr7 a b => do
+    let a' ← checkNoEscape a .field
+    let b' ← checkNoEscape b .field
+    pure (Typed.Term.u8ChainRotr7 (.tuple #[.field, .field, .field]) false a' b')
+  | .u8ChainRotr4 a b => do
+    let a' ← checkNoEscape a .field
+    let b' ← checkNoEscape b .field
+    pure (Typed.Term.u8ChainRotr4 (.tuple #[.field, .field, .field]) false a' b')
   | .u8Sub a b => do
     let a' ← checkNoEscape a .field
     let b' ← checkNoEscape b .field
@@ -905,6 +913,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
   | .u8Xor τ e a b => do pure (.u8Xor (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Add τ e a b => do pure (.u8Add (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Mul τ e a b => do pure (.u8Mul (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
+  | .u8ChainRotr7 τ e a b => do pure (.u8ChainRotr7 (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
+  | .u8ChainRotr4 τ e a b => do pure (.u8ChainRotr4 (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Sub τ e a b => do pure (.u8Sub (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8And τ e a b => do pure (.u8And (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Or τ e a b => do pure (.u8Or (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))

Ix/Aiur/Compiler/Concretize.lean

@@ -328,6 +328,12 @@ def termToConcrete
   | .u8Mul τ e a b => do
       pure (.u8Mul (← typToConcrete mono τ) e
                    (← termToConcrete mono a) (← termToConcrete mono b))
+  | .u8ChainRotr7 τ e a b => do
+      pure (.u8ChainRotr7 (← typToConcrete mono τ) e
+                   (← termToConcrete mono a) (← termToConcrete mono b))
+  | .u8ChainRotr4 τ e a b => do
+      pure (.u8ChainRotr4 (← typToConcrete mono τ) e
+                   (← termToConcrete mono a) (← termToConcrete mono b))
   | .u8Sub τ e a b => do
       pure (.u8Sub (← typToConcrete mono τ) e
                    (← termToConcrete mono a) (← termToConcrete mono b))
@@ -526,6 +532,10 @@ def rewriteTypedTerm (decls : Typed.Decls)
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
   | .u8Mul τ e a b => .u8Mul (rewriteTyp subst mono τ) e
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
+  | .u8ChainRotr7 τ e a b => .u8ChainRotr7 (rewriteTyp subst mono τ) e
+      (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
+  | .u8ChainRotr4 τ e a b => .u8ChainRotr4 (rewriteTyp subst mono τ) e
+      (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
   | .u8Sub τ e a b => .u8Sub (rewriteTyp subst mono τ) e
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
   | .u8And τ e a b => .u8And (rewriteTyp subst mono τ) e
@@ -603,6 +613,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
     args.attach.foldl (fun s ⟨a, _⟩ => collectInTypedTerm s a) seen
   | .add τ _ a b | .sub τ _ a b | .mul τ _ a b
   | .u8Xor τ _ a b | .u8Add τ _ a b | .u8Mul τ _ a b | .u8Sub τ _ a b
+  | .u8ChainRotr7 τ _ a b | .u8ChainRotr4 τ _ a b
   | .u8And τ _ a b | .u8Or τ _ a b
   | .u8LessThan τ _ a b | .u32LessThan τ _ a b =>
     collectInTypedTerm (collectInTypedTerm (collectInTyp seen τ) a) b
@@ -665,6 +676,7 @@ def collectCalls (decls : Typed.Decls)
     bs.attach.foldl (fun s ⟨(_, b), _⟩ => collectCalls decls s b) seen
   | .add _ _ a b | .sub _ _ a b | .mul _ _ a b
   | .u8Xor _ _ a b | .u8Add _ _ a b | .u8Mul _ _ a b | .u8Sub _ _ a b
+  | .u8ChainRotr7 _ _ a b | .u8ChainRotr4 _ _ a b
   | .u8And _ _ a b | .u8Or _ _ a b
   | .u8LessThan _ _ a b | .u32LessThan _ _ a b =>
     collectCalls decls (collectCalls decls seen a) b
@@ -756,6 +768,10 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
       (substInTypedTerm subst a) (substInTypedTerm subst b)
   | .u8Mul τ e a b => .u8Mul (Typ.instantiate subst τ) e
       (substInTypedTerm subst a) (substInTypedTerm subst b)
+  | .u8ChainRotr7 τ e a b => .u8ChainRotr7 (Typ.instantiate subst τ) e
+      (substInTypedTerm subst a) (substInTypedTerm subst b)
+  | .u8ChainRotr4 τ e a b => .u8ChainRotr4 (Typ.instantiate subst τ) e
+      (substInTypedTerm subst a) (substInTypedTerm subst b)
   | .u8Sub τ e a b => .u8Sub (Typ.instantiate subst τ) e
       (substInTypedTerm subst a) (substInTypedTerm subst b)
   | .u8And τ e a b => .u8And (Typ.instantiate subst τ) e

Ix/Aiur/Compiler/Layout.lean

@@ -185,7 +185,17 @@ def opLayout : Bytecode.Op → LayoutM Unit
   | .u8BitDecomposition _ => do pushDegrees $ .replicate 8 1; bumpAuxiliaries 8; bumpLookups
   | .u8ShiftLeft _ | .u8ShiftRight _ | .u8Xor .. | .u8And .. | .u8Or .. => do
     pushDegree 1; bumpAuxiliaries; bumpLookups
-  | .u8Add .. | .u8Mul .. | .u8Sub .. => do pushDegrees #[1, 1]; bumpAuxiliaries 2; bumpLookups
+  | .u8Add a b | .u8Sub a b => do
+    -- Low byte `z` is the only auxiliary; the carry/borrow is a compound
+    -- expression of degree `max(deg a, deg b, 1)` (add: `(a+b-z)/256`,
+    -- sub: `(z+b-a)/256`).
+    let aDegree ← getDegree a
+    let bDegree ← getDegree b
+    pushDegree 1
+    pushDegree ((aDegree.max bDegree).max 1)
+    bumpAuxiliaries; bumpLookups
+  | .u8Mul .. => do pushDegrees #[1, 1]; bumpAuxiliaries 2; bumpLookups
+  | .u8ChainRotr7 .. | .u8ChainRotr4 .. => do pushDegrees #[1, 1, 1]; bumpAuxiliaries 3; bumpLookups
   | .u8LessThan .. => do pushDegree 1; bumpAuxiliaries; bumpLookups
   | .u32LessThan .. => do pushDegree 1; bumpAuxiliaries 12; bumpLookups 6
   | .debug .. => pure ()

Ix/Aiur/Compiler/Lower.lean

@@ -286,6 +286,12 @@ def toIndex
   | .u8Mul _ _ i j => do
     let i ← expectIdx layoutMap bindings i; let j ← expectIdx layoutMap bindings j
     pushOp (.u8Mul i j) 2
+  | .u8ChainRotr7 _ _ i j => do
+    let i ← expectIdx layoutMap bindings i; let j ← expectIdx layoutMap bindings j
+    pushOp (.u8ChainRotr7 i j) 3
+  | .u8ChainRotr4 _ _ i j => do
+    let i ← expectIdx layoutMap bindings i; let j ← expectIdx layoutMap bindings j
+    pushOp (.u8ChainRotr4 i j) 3
   | .u8Sub _ _ i j => do
     let i ← expectIdx layoutMap bindings i; let j ← expectIdx layoutMap bindings j
     pushOp (.u8Sub i j) 2

Ix/Aiur/Compiler/Match.lean

@@ -381,6 +381,8 @@ def typedToSimple : Term → Simple.Term
   | .u8Xor τ e a b => .u8Xor τ e (typedToSimple a) (typedToSimple b)
   | .u8Add τ e a b => .u8Add τ e (typedToSimple a) (typedToSimple b)
   | .u8Mul τ e a b => .u8Mul τ e (typedToSimple a) (typedToSimple b)
+  | .u8ChainRotr7 τ e a b => .u8ChainRotr7 τ e (typedToSimple a) (typedToSimple b)
+  | .u8ChainRotr4 τ e a b => .u8ChainRotr4 τ e (typedToSimple a) (typedToSimple b)
   | .u8Sub τ e a b => .u8Sub τ e (typedToSimple a) (typedToSimple b)
   | .u8And τ e a b => .u8And τ e (typedToSimple a) (typedToSimple b)
   | .u8Or τ e a b => .u8Or τ e (typedToSimple a) (typedToSimple b)

Ix/Aiur/Goldilocks.lean

@@ -64,6 +64,16 @@ def G.u8Mul (a b : G) : G × G :=
 def G.u8Sub (a b : G) : G × G :=
   (G.ofNat ((a.n + 256 - b.n) % 256), if a.n < b.n then 1 else 0)
 
+/-- Chainable partial for a right-rotation by 7 bits over little-endian bytes.
+Returns `(a>>7 + b<<1, b>>7, a<<1)` (shifts mod 256). -/
+def G.u8ChainRotr7 (a b : G) : G × G × G :=
+  (G.ofNat (a.n / 128 + (b.n * 2) % 256), G.ofNat (b.n / 128), G.ofNat ((a.n * 2) % 256))
+
+/-- Chainable partial for a right-rotation by 4 bits over little-endian bytes.
+Returns `(a>>4 + b<<4, b>>4, a<<4)` (shifts mod 256). -/
+def G.u8ChainRotr4 (a b : G) : G × G × G :=
+  (G.ofNat (a.n / 16 + (b.n * 16) % 256), G.ofNat (b.n / 16), G.ofNat ((a.n * 16) % 256))
+
 def G.u8ShiftLeft  (a : G) : G := G.ofNat ((a.n * 2) % 256)
 def G.u8ShiftRight (a : G) : G := G.ofNat (a.n / 2)
 

Ix/Aiur/Interpret.lean

@@ -358,6 +358,24 @@ partial def interp (decls : Decls) (bindings : Bindings) : Term → InterpM Valu
           let hi : Value := .field (G.ofUInt8 (x / 256).toUInt8)
           return .tuple #[lo, hi]
       | _, _ => throwErr "u8Mul: expected field values"
+  | .u8ChainRotr7 t1 t2 => do
+      match (← interp decls bindings t1), (← interp decls bindings t2) with
+      | .field a, .field b =>
+          let i := a.val.toUInt8
+          let j := b.val.toUInt8
+          return .tuple #[.field (G.ofUInt8 ((i >>> 7) + (j <<< 1))),
+                          .field (G.ofUInt8 (j >>> 7)),
+                          .field (G.ofUInt8 (i <<< 1))]
+      | _, _ => throwErr "u8ChainRotr7: expected field values"
+  | .u8ChainRotr4 t1 t2 => do
+      match (← interp decls bindings t1), (← interp decls bindings t2) with
+      | .field a, .field b =>
+          let i := a.val.toUInt8
+          let j := b.val.toUInt8
+          return .tuple #[.field (G.ofUInt8 ((i >>> 4) + (j <<< 4))),
+                          .field (G.ofUInt8 (j >>> 4)),
+                          .field (G.ofUInt8 (i <<< 4))]
+      | _, _ => throwErr "u8ChainRotr4: expected field values"
   | .u8Sub t1 t2 => do
       match (← interp decls bindings t1), (← interp decls bindings t2) with
       | .field a, .field b =>

Ix/Aiur/Meta.lean

@@ -174,6 +174,8 @@ syntax "u8_shift_right" "(" aiur_trm ")"                                     : a
 syntax "u8_xor" "(" aiur_trm ", " aiur_trm ")"                              : aiur_trm
 syntax "u8_add" "(" aiur_trm ", " aiur_trm ")"                              : aiur_trm
 syntax "u8_mul" "(" aiur_trm ", " aiur_trm ")"                              : aiur_trm
+syntax "u8_chain_rotr7" "(" aiur_trm ", " aiur_trm ")"                      : aiur_trm
+syntax "u8_chain_rotr4" "(" aiur_trm ", " aiur_trm ")"                      : aiur_trm
 syntax "u8_sub" "(" aiur_trm ", " aiur_trm ")"                              : aiur_trm
 syntax "u8_and" "(" aiur_trm ", " aiur_trm ")"                              : aiur_trm
 syntax "u8_or" "(" aiur_trm ", " aiur_trm ")"                               : aiur_trm
@@ -295,6 +297,10 @@ partial def elabTrm : ElabStxCat `aiur_trm
     mkAppM ``Source.Term.u8Add #[← elabTrm i, ← elabTrm j]
   | `(aiur_trm| u8_mul($i:aiur_trm, $j:aiur_trm)) => do
     mkAppM ``Source.Term.u8Mul #[← elabTrm i, ← elabTrm j]
+  | `(aiur_trm| u8_chain_rotr7($i:aiur_trm, $j:aiur_trm)) => do
+    mkAppM ``Source.Term.u8ChainRotr7 #[← elabTrm i, ← elabTrm j]
+  | `(aiur_trm| u8_chain_rotr4($i:aiur_trm, $j:aiur_trm)) => do
+    mkAppM ``Source.Term.u8ChainRotr4 #[← elabTrm i, ← elabTrm j]
   | `(aiur_trm| u8_sub($i:aiur_trm, $j:aiur_trm)) => do
     mkAppM ``Source.Term.u8Sub #[← elabTrm i, ← elabTrm j]
   | `(aiur_trm| u8_and($i:aiur_trm, $j:aiur_trm)) => do
@@ -487,6 +493,14 @@ where
       let i ← replaceToken old new i
       let j ← replaceToken old new j
       `(aiur_trm| u8_mul($i, $j))
+    | `(aiur_trm| u8_chain_rotr7($i:aiur_trm, $j:aiur_trm)) => do
+      let i ← replaceToken old new i
+      let j ← replaceToken old new j
+      `(aiur_trm| u8_chain_rotr7($i, $j))
+    | `(aiur_trm| u8_chain_rotr4($i:aiur_trm, $j:aiur_trm)) => do
+      let i ← replaceToken old new i
+      let j ← replaceToken old new j
+      `(aiur_trm| u8_chain_rotr4($i, $j))
     | `(aiur_trm| u8_sub($i:aiur_trm, $j:aiur_trm)) => do
       let i ← replaceToken old new i
       let j ← replaceToken old new j

Ix/Aiur/Semantics/BytecodeEval.lean

@@ -252,6 +252,18 @@ def evalOp (t : Bytecode.Toplevel) (fuel : Nat) (op : Op) (st : EvalState) :
     let prod := x.val.toUInt8.toNat * y.val.toUInt8.toNat
     let st1 := pushMap st (G.ofUInt8 prod.toUInt8)
     pure (pushMap st1 (G.ofUInt8 (prod / 256).toUInt8))
+  | .u8ChainRotr7 a b => do
+    let x ← readIdx st a; let y ← readIdx st b
+    let i := x.val.toUInt8; let j := y.val.toUInt8
+    let st1 := pushMap st (G.ofUInt8 ((i >>> 7) + (j <<< 1)))
+    let st2 := pushMap st1 (G.ofUInt8 (j >>> 7))
+    pure (pushMap st2 (G.ofUInt8 (i <<< 1)))
+  | .u8ChainRotr4 a b => do
+    let x ← readIdx st a; let y ← readIdx st b
+    let i := x.val.toUInt8; let j := y.val.toUInt8
+    let st1 := pushMap st (G.ofUInt8 ((i >>> 4) + (j <<< 4)))
+    let st2 := pushMap st1 (G.ofUInt8 (j >>> 4))
+    pure (pushMap st2 (G.ofUInt8 (i <<< 4)))
   | .u8Sub a b => do
     let x ← readIdx st a; let y ← readIdx st b
     let i := x.val.toUInt8; let j := y.val.toUInt8

Ix/Aiur/Semantics/SourceEval.lean

@@ -390,6 +390,24 @@ def interp (decls : Decls) (fuel : Nat) (bindings : Bindings)
                    .field (G.ofUInt8 (x / 256).toUInt8)])
         (interp decls fuel bindings t1 st)
         (fun st1 => interp decls fuel bindings t2 st1)
+  | .u8ChainRotr7 t1 t2 =>
+      combineFieldsResult
+        (fun a b =>
+          let i := a.val.toUInt8; let j := b.val.toUInt8
+          .tuple #[.field (G.ofUInt8 ((i >>> 7) + (j <<< 1))),
+                   .field (G.ofUInt8 (j >>> 7)),
+                   .field (G.ofUInt8 (i <<< 1))])
+        (interp decls fuel bindings t1 st)
+        (fun st1 => interp decls fuel bindings t2 st1)
+  | .u8ChainRotr4 t1 t2 =>
+      combineFieldsResult
+        (fun a b =>
+          let i := a.val.toUInt8; let j := b.val.toUInt8
+          .tuple #[.field (G.ofUInt8 ((i >>> 4) + (j <<< 4))),
+                   .field (G.ofUInt8 (j >>> 4)),
+                   .field (G.ofUInt8 (i <<< 4))])
+        (interp decls fuel bindings t1 st)
+        (fun st1 => interp decls fuel bindings t2 st1)
   | .u8Sub t1 t2 =>
       combineFieldsResult
         (fun a b =>