Fused sqrt and invsqrt

[mratsim]

I've noticed that you integrated hash-to-G2 from the H2C draft and are computing sqrt and invsqrt in a separate manner here:

core/c/fp2.c

Lines 456 to 458 in bed0047

	FP_YYY_sqrt(&(w->a),&w2,&hint); // a=sqrt(w2)
	FP_YYY_inv(&w3,&w2,&hint); // w3=1/w2
	FP_YYY_mul(&w3,&w3,&(w->a)); // w3=1/sqrt(w2)

Instead you can compute sqrt and invsqrt in a fused manner by modifying your Tonelli Shanks algorithm by not multiplying Line 799

core/c/fp.c

Lines 782 to 817 in bed0047

	void FP_YYY_sqrt(FP_YYY *r, FP_YYY *a, FP_YYY *h)
	{
	int i,j,k,u,e=PM1D2_YYY;
	FP_YYY v,g,t,b;
	BIG_XXX m;
	if (h==NULL)
	FP_YYY_progen(&g,a);
	else
	FP_YYY_copy(&g,h);
	BIG_XXX_rcopy(m,ROI_YYY);
	FP_YYY_nres(&v,m);
	FP_YYY_sqr(&t,&g);
	FP_YYY_mul(&t,&t,a);
	FP_YYY_mul(r,&g,a);
	FP_YYY_copy(&b,&t);
	for (k=e;k>1;k--)
	{
	for (j=1;j<k-1;j++)
	FP_YYY_sqr(&b,&b);
	u=1-FP_YYY_isunity(&b);
	FP_YYY_mul(&g,r,&v);
	FP_YYY_cmove(r,&g,u);
	FP_YYY_sqr(&v,&v);
	FP_YYY_mul(&g,&t,&v);
	FP_YYY_cmove(&t,&g,u);
	FP_YYY_copy(&b,&t);
	}
	// always return +ve square root
	k=FP_YYY_sign(r);
	FP_YYY_neg(&v,r); FP_YYY_norm(&v);
	FP_YYY_cmove(r,&v,k);
	}

Then r will hold invsqrt(a) at the end of the Tonelli Shanks, and you can multiply that by a and return both sqrt and invsqrt in a fused manner.

My own implementation (in Nim):
https://github.com/mratsim/constantine/blob/0effd66d/constantine/arithmetic/finite_fields_square_root.nim#L147-L180

func sqrt_invsqrt_tonelli_shanks[C](
       sqrt, invsqrt: var Fp[C],
       a, a_pre_exp: Fp[C]) =
  ## Compute the square_root and inverse_square_root
  ## of `a` via constant-time Tonelli-Shanks
  ##
  ## a_pre_exp is a precomputation a^((p-1-2^e)/(2*2^e))
  template z: untyped = a_pre_exp
  template r: untyped = invsqrt
  var t {.noInit.}: Fp[C]
  const e = C.tonelliShanks(twoAdicity)

  t.square(z)
  t *= a
  r = z
  var b = t
  var root = C.tonelliShanks(root_of_unity)

  var buf {.noInit.}: Fp[C]

  for i in countdown(e, 2, 1):
    for j in 1 .. i-2:
      b.square()

    let bNotOne = not b.isOne()
    buf.prod(r, root)
    r.ccopy(buf, bNotOne)
    root.square()
    buf.prod(t, root)
    t.ccopy(buf, bNotOne)
    b = t

  sqrt.prod(invsqrt, a)

Sage script for BLS12-377
https://github.com/mratsim/constantine/blob/0effd66d/sage/square_root_bls12_377.sage

def sqrt_inv_sqrt_tonelli_shanks_impl(a, a_pre_exp, s, e, root_of_unity):
    ## Square root and inverse square root for any `a` in a field of prime characteristic p
    ##
    ## a_pre_exp = a^((q-1-2^e)/(2*2^e))
    ## with
    ##  s and e, precomputed values
    ##  such as q == s * 2^e + 1

    # Implementation
    # 1/√a * a = √a
    # Notice that in Tonelli Shanks, the result `r` is bootstrapped by "z*a"
    # We bootstrap it instead by just z to get invsqrt for free

    z = a_pre_exp
    t = z*z*a
    r = z
    b = t
    root = root_of_unity
    for i in range(e, 1, -1):   # e .. 2
        for j in range(1, i-1): # 1 .. i-2
            b *= b
        doCopy = b != 1
        r = ccopy(r, r * root, doCopy)
        root *= root
        t = ccopy(t, t * root, doCopy)
        b = t
    return r*a, r

[mcarrickscott]

That's very neat! In fact I wasn't computing the sqrt and the inverse completely separately, as they both share the same "hint" value which does most of the heavy lifting. But what you suggest would be even faster.