koalabear mixed vecops

ingonyama-zk · Jan 1, 2025 · 1dc168e · 1dc168e
1 parent 8e2c1e3
commit 1dc168e
Show file tree

Hide file tree

Showing 10 changed files with 91 additions and 3 deletions.
diff --git a/wrappers/golang/fields/koalabear/extension/extension_field.go b/wrappers/golang/fields/koalabear/extension/extension_field.go
@@ -168,6 +168,18 @@ func (f ExtensionField) Sqr() ExtensionField {
 	return res
 }
 
+func (f ExtensionField) Pow(exp int) ExtensionField {
+	var res ExtensionField
+
+	cF := (*C.scalar_t)(unsafe.Pointer(&f))
+	cExp := (C.int)(exp)
+	cRes := (*C.scalar_t)(unsafe.Pointer(&res))
+
+	C.koalabear_extension_pow(cF, cExp, cRes)
+
+	return res
+}
+
 func convertScalarsMontgomery(scalars core.HostOrDeviceSlice, isInto bool) runtime.EIcicleError {
 	defaultCfg := core.DefaultVecOpsConfig()
 	cValues, _, _, cCfg, cSize := core.VecOpCheck(scalars, scalars, scalars, &defaultCfg)

diff --git a/wrappers/golang/fields/koalabear/extension/include/scalar_field.h b/wrappers/golang/fields/koalabear/extension/include/scalar_field.h
@@ -16,6 +16,7 @@ void koalabear_extension_add(const scalar_t* a, const scalar_t* b, scalar_t* res
 void koalabear_extension_sub(const scalar_t* a, const scalar_t* b, scalar_t* result);
 void koalabear_extension_mul(const scalar_t* a, const scalar_t* b, scalar_t* result);
 void koalabear_extension_inv(const scalar_t* a, scalar_t* result);
+void koalabear_extension_pow(const scalar_t* a, int exp, scalar_t* result);
 
 #ifdef __cplusplus
 }

diff --git a/wrappers/golang/fields/koalabear/extension/vecOps/include/vec_ops.h b/wrappers/golang/fields/koalabear/extension/vecOps/include/vec_ops.h
@@ -43,6 +43,14 @@ int koalabear_extension_matrix_transpose(
   scalar_t* mat_out
 );
 
+int koalabear_extension_vector_mixed_mul(
+  scalar_t* vec_a,
+  scalar_t* vec_b,
+  int n,
+  VecOpsConfig* config,
+  scalar_t* result
+);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/wrappers/golang/fields/koalabear/extension/vecOps/vec_ops.go b/wrappers/golang/fields/koalabear/extension/vecOps/vec_ops.go
@@ -42,3 +42,20 @@ func TransposeMatrix(in, out core.HostOrDeviceSlice, columnSize, rowSize int, co
 	err := (C.koalabear_extension_matrix_transpose(cIn, cRowSize, cColumnSize, cConfig, cOut))
 	return runtime.EIcicleError(err)
 }
+
+func MixedVecOp(a, b, out core.HostOrDeviceSlice, config core.VecOpsConfig, op core.VecOps) (ret runtime.EIcicleError) {
+	aPointer, bPointer, outPointer, cfgPointer, size := core.VecOpCheck(a, b, out, &config)
+
+	cA := (*C.scalar_t)(aPointer)
+	cB := (*C.scalar_t)(bPointer)
+	cOut := (*C.scalar_t)(outPointer)
+	cConfig := (*C.VecOpsConfig)(cfgPointer)
+	cSize := (C.int)(size)
+
+	switch op {
+	case core.Mul:
+		ret = (runtime.EIcicleError)(C.koalabear_extension_vector_mixed_mul(cA, cB, cSize, cConfig, cOut))
+	}
+
+	return ret
+}
diff --git a/wrappers/golang/fields/koalabear/include/scalar_field.h b/wrappers/golang/fields/koalabear/include/scalar_field.h
@@ -16,6 +16,7 @@ void koalabear_add(const scalar_t* a, const scalar_t* b, scalar_t* result);
 void koalabear_sub(const scalar_t* a, const scalar_t* b, scalar_t* result);
 void koalabear_mul(const scalar_t* a, const scalar_t* b, scalar_t* result);
 void koalabear_inv(const scalar_t* a, scalar_t* result);
+void koalabear_pow(const scalar_t* a, int exp, scalar_t* result);
 
 #ifdef __cplusplus
 }

diff --git a/wrappers/golang/fields/koalabear/scalar_field.go b/wrappers/golang/fields/koalabear/scalar_field.go
@@ -168,6 +168,18 @@ func (f ScalarField) Sqr() ScalarField {
 	return res
 }
 
+func (f ScalarField) Pow(exp int) ScalarField {
+	var res ScalarField
+
+	cF := (*C.scalar_t)(unsafe.Pointer(&f))
+	cExp := (C.int)(exp)
+	cRes := (*C.scalar_t)(unsafe.Pointer(&res))
+
+	C.koalabear_pow(cF, cExp, cRes)
+
+	return res
+}
+
 func convertScalarsMontgomery(scalars core.HostOrDeviceSlice, isInto bool) runtime.EIcicleError {
 	defaultCfg := core.DefaultVecOpsConfig()
 	cValues, _, _, cCfg, cSize := core.VecOpCheck(scalars, scalars, scalars, &defaultCfg)

diff --git a/wrappers/golang/fields/koalabear/tests/extension_field_test.go b/wrappers/golang/fields/koalabear/tests/extension_field_test.go
@@ -119,6 +119,11 @@ func testExtensionFieldArithmetic(suite *suite.Suite) {
 
 	suite.Equal(square, mul, "Square and multiplication do not yield the same value")
 
+	pow4 := scalarA.Pow(4)
+	mulBySelf := mul.Mul(&mul)
+
+	suite.Equal(pow4, mulBySelf, "Square and multiplication do not yield the same value")
+
 	inv := scalarA.Inv()
 
 	one := scalarA.Mul(&inv)

diff --git a/wrappers/golang/fields/koalabear/tests/extension_vec_ops_test.go b/wrappers/golang/fields/koalabear/tests/extension_vec_ops_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/ingonyama-zk/icicle/v3/wrappers/golang/core"
+	koalabear "github.com/ingonyama-zk/icicle/v3/wrappers/golang/fields/koalabear"
 	koalabear_extension "github.com/ingonyama-zk/icicle/v3/wrappers/golang/fields/koalabear/extension"
 	"github.com/ingonyama-zk/icicle/v3/wrappers/golang/fields/koalabear/extension/vecOps"
 	"github.com/stretchr/testify/suite"
@@ -64,13 +65,31 @@ func testKoalabear_extensionTranspose(suite *suite.Suite) {
 	suite.Equal(matrix, output)
 }
 
+func testKoalabear_extensionMixedVecOps(suite *suite.Suite) {
+	testSize := 1 << 14
+
+	a := koalabear_extension.GenerateScalars(testSize)
+	var scalar koalabear.ScalarField
+	scalar.One()
+	ones := core.HostSliceWithValue(scalar, testSize)
+
+	out := make(core.HostSlice[koalabear_extension.ExtensionField], testSize)
+
+	cfg := core.DefaultVecOpsConfig()
+
+	vecOps.MixedVecOp(a, ones, out, cfg, core.Mul)
+
+	suite.Equal(a, out)
+}
+
 type Koalabear_extensionVecOpsTestSuite struct {
 	suite.Suite
 }
 
 func (s *Koalabear_extensionVecOpsTestSuite) TestKoalabear_extensionVecOps() {
 	s.Run("TestKoalabear_extensionVecOps", testWrapper(&s.Suite, testKoalabear_extensionVecOps))
 	s.Run("TestKoalabear_extensionTranspose", testWrapper(&s.Suite, testKoalabear_extensionTranspose))
+	s.Run("TestKoalabear_extensionMixedVecOps", testWrapper(&s.Suite, testKoalabear_extensionMixedVecOps))
 }
 
 func TestSuiteKoalabear_extensionVecOps(t *testing.T) {

diff --git a/wrappers/golang/fields/koalabear/tests/scalar_field_test.go b/wrappers/golang/fields/koalabear/tests/scalar_field_test.go
@@ -119,6 +119,11 @@ func testScalarFieldArithmetic(suite *suite.Suite) {
 
 	suite.Equal(square, mul, "Square and multiplication do not yield the same value")
 
+	pow4 := scalarA.Pow(4)
+	mulBySelf := mul.Mul(&mul)
+
+	suite.Equal(pow4, mulBySelf, "Square and multiplication do not yield the same value")
+
 	inv := scalarA.Inv()
 
 	one := scalarA.Mul(&inv)

diff --git a/wrappers/rust/icicle-fields/icicle-koalabear/src/vec_ops/mod.rs b/wrappers/rust/icicle-fields/icicle-koalabear/src/vec_ops/mod.rs
@@ -1,24 +1,32 @@
 use crate::field::{ExtensionCfg, ExtensionField, ScalarCfg, ScalarField};
 
-use icicle_core::impl_vec_ops_field;
-use icicle_core::vec_ops::{VecOps, VecOpsConfig};
+use icicle_core::vec_ops::{MixedVecOps, VecOps, VecOpsConfig};
+use icicle_core::{impl_vec_ops_field, impl_vec_ops_mixed_field};
 use icicle_runtime::errors::eIcicleError;
 use icicle_runtime::memory::HostOrDeviceSlice;
 
 impl_vec_ops_field!("koalabear", koalabear, ScalarField, ScalarCfg);
 impl_vec_ops_field!("koalabear_extension", koalabear_extension, ExtensionField, ExtensionCfg);
+impl_vec_ops_mixed_field!(
+    "koalabear_extension",
+    koalabear_mixed,
+    ExtensionField,
+    ScalarField,
+    ExtensionCfg
+);
 
 #[cfg(test)]
 pub(crate) mod tests {
     use crate::field::{ExtensionField, ScalarField};
-    use icicle_core::impl_vec_ops_tests;
     use icicle_core::vec_ops::tests::*;
+    use icicle_core::{impl_mixed_vec_ops_tests, impl_vec_ops_tests};
 
     impl_vec_ops_tests!(ScalarField);
 
     mod extension {
         use super::*;
 
         impl_vec_ops_tests!(ExtensionField);
+        impl_mixed_vec_ops_tests!(ExtensionField, ScalarField);
     }
 }