From f613ab2c257ad42038ce2f78a8eb858cf840a3a5 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 9 Oct 2019 11:37:22 +0300
Subject: [PATCH] Implement matrix multiplication with pure Go

Set suitable backend for GF127 arithmetic for Concat(), Sum() etc.
---
 gogf127/gogf127.go |  9 +++++----
 tz/cpuid_x86.go    |  2 +-
 tz/sl2.go          | 35 +++++++++++++++++++++++++++++++----
 3 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/gogf127/gogf127.go b/gogf127/gogf127.go
index cff06a3..8ba33e9 100644
--- a/gogf127/gogf127.go
+++ b/gogf127/gogf127.go
@@ -160,23 +160,24 @@ func Add(a, b, c *GF127) {
 }
 
 // Mul sets c to a*b.
+// TODO make it work in-place without allocations
 // TODO optimization: no need to perform shift by i every time, cache results
 func Mul(a, b, c *GF127) {
-	c[0] = 0
-	c[1] = 0
+	r := new(GF127)
 	d := new(GF127)
 	for i := uint(0); i < 64; i++ {
 		if b[0]&(1<<i) != 0 {
 			shl(i, a, d)
-			Add(c, d, c)
+			Add(r, d, r)
 		}
 	}
 	for i := uint(0); i < 63; i++ {
 		if b[1]&(1<<i) != 0 {
 			shl(i+64, a, d)
-			Add(c, d, c)
+			Add(r, d, r)
 		}
 	}
+	*c = *r
 }
 
 // shl performs left shift by consecutive multiplications by 2.
diff --git a/tz/cpuid_x86.go b/tz/cpuid_x86.go
index 60d0b1b..1d7630e 100644
--- a/tz/cpuid_x86.go
+++ b/tz/cpuid_x86.go
@@ -36,7 +36,7 @@ const (
 	bitAVX2    = 1 << 5
 )
 
-func init() {
+func setFeatures() {
 	maxID, _, _, _ := cpuid(0, 0)
 	if maxID < 1 {
 		return
diff --git a/tz/sl2.go b/tz/sl2.go
index ea56db0..231de4e 100644
--- a/tz/sl2.go
+++ b/tz/sl2.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 
 	"github.com/nspcc-dev/tzhash/gf127"
+	"github.com/nspcc-dev/tzhash/gogf127"
 )
 
 type sl2 [2][2]gf127.GF127
@@ -13,6 +14,18 @@ var id = sl2{
 	{gf127.GF127{0, 0}, gf127.GF127{1, 0}},
 }
 
+var mul func(a, b, c *sl2, x *[4]gf127.GF127)
+
+func init() {
+	setFeatures()
+
+	if hasAVX {
+		mul = mulSL2AVX
+	} else {
+		mul = mulSL2Pure
+	}
+}
+
 func (c *sl2) MarshalBinary() (data []byte, err error) {
 	s := c.ByteArray()
 	return s[:], nil
@@ -79,8 +92,7 @@ func (c *sl2) mulStrassen(a, b *sl2, x *[8]gf127.GF127) *sl2 {
 	return c
 }
 
-func (c *sl2) mul(a, b *sl2, x *[4]gf127.GF127) *sl2 {
-	// naive implementation
+func mulSL2AVX(a, b, c *sl2, x *[4]gf127.GF127) {
 	gf127.Mul(&a[0][0], &b[0][0], &x[0])
 	gf127.Mul(&a[0][0], &b[0][1], &x[1])
 	gf127.Mul(&a[1][0], &b[0][0], &x[2])
@@ -94,8 +106,22 @@ func (c *sl2) mul(a, b *sl2, x *[4]gf127.GF127) *sl2 {
 	gf127.Add(&c[1][0], &x[2], &c[1][0])
 	gf127.Mul(&a[1][1], &b[1][1], &c[1][1])
 	gf127.Add(&c[1][1], &x[3], &c[1][1])
+}
 
-	return c
+func mulSL2Pure(a, b, c *sl2, x *[4]gf127.GF127) {
+	gogf127.Mul((*gogf127.GF127)(&a[0][0]), (*gogf127.GF127)(&b[0][0]), (*gogf127.GF127)(&x[0]))
+	gogf127.Mul((*gogf127.GF127)(&a[0][0]), (*gogf127.GF127)(&b[0][1]), (*gogf127.GF127)(&x[1]))
+	gogf127.Mul((*gogf127.GF127)(&a[1][0]), (*gogf127.GF127)(&b[0][0]), (*gogf127.GF127)(&x[2]))
+	gogf127.Mul((*gogf127.GF127)(&a[1][0]), (*gogf127.GF127)(&b[0][1]), (*gogf127.GF127)(&x[3]))
+
+	gogf127.Mul((*gogf127.GF127)(&a[0][1]), (*gogf127.GF127)(&b[1][0]), (*gogf127.GF127)(&c[0][0]))
+	gogf127.Add((*gogf127.GF127)(&c[0][0]), (*gogf127.GF127)(&x[0]), (*gogf127.GF127)(&c[0][0]))
+	gogf127.Mul((*gogf127.GF127)(&a[0][1]), (*gogf127.GF127)(&b[1][1]), (*gogf127.GF127)(&c[0][1]))
+	gogf127.Add((*gogf127.GF127)(&c[0][1]), (*gogf127.GF127)(&x[1]), (*gogf127.GF127)(&c[0][1]))
+	gogf127.Mul((*gogf127.GF127)(&a[1][1]), (*gogf127.GF127)(&b[1][0]), (*gogf127.GF127)(&c[1][0]))
+	gogf127.Add((*gogf127.GF127)(&c[1][0]), (*gogf127.GF127)(&x[2]), (*gogf127.GF127)(&c[1][0]))
+	gogf127.Mul((*gogf127.GF127)(&a[1][1]), (*gogf127.GF127)(&b[1][1]), (*gogf127.GF127)(&c[1][1]))
+	gogf127.Add((*gogf127.GF127)(&c[1][1]), (*gogf127.GF127)(&x[3]), (*gogf127.GF127)(&c[1][1]))
 }
 
 func (c *sl2) MulA() *sl2 {
@@ -129,7 +155,8 @@ func (c *sl2) MulB() *sl2 {
 }
 
 func (c *sl2) Mul(a, b *sl2) *sl2 {
-	return c.mul(a, b, new([4]gf127.GF127))
+	mul(a, b, c, new([4]gf127.GF127))
+	return c
 }
 
 // Inv returns inverse of a in GL_2(GF(2^127))