diff --git a/c/fft.cpp b/c/fft.cpp index c557ae0..cdb9f0f 100644 --- a/c/fft.cpp +++ b/c/fft.cpp @@ -114,6 +114,36 @@ FFT::FFT(u_int64_t maxDomainSize, uint32_t _nThreads) mpz_clear(m_aux); } +template +void FFT::higherRootOfUnity(Element &r, u_int32_t extraPow) { + mpz_t m_q, m_aux, m_nqr; + + mpz_init(m_q); + mpz_init(m_aux); + mpz_init(m_nqr); + + f.toMpz(m_aux, f.negOne()); + mpz_add_ui(m_q, m_aux, 1); + + // (q-1) / 2^(s+extraPow); the primitive root exists iff the division + // is exact, i.e. s+extraPow is within the field's 2-adicity + if (mpz_scan1(m_aux, 0) < s + extraPow) { + mpz_clear(m_q); + mpz_clear(m_aux); + mpz_clear(m_nqr); + throw std::range_error("Root order exceeds the field's 2-adicity"); + } + mpz_fdiv_q_2exp(m_aux, m_aux, s + extraPow); + + f.toMpz(m_nqr, nqr); + mpz_powm(m_aux, m_nqr, m_aux, m_q); + f.fromMpz(r, m_aux); + + mpz_clear(m_q); + mpz_clear(m_aux); + mpz_clear(m_nqr); +} + template FFT::~FFT() { delete[] roots; @@ -198,6 +228,60 @@ void FFT::fft(Element *a, u_int64_t n) { } } +template +void FFT::fftDITRevToNat(Element *a, u_int64_t n) { + u_int64_t domainPow = log2(n); + assert(((u_int64_t)1 << domainPow) == n); + + for (u_int32_t s=1; s<=domainPow; s++) { + u_int64_t m = 1 << s; + u_int64_t mdiv2 = m >> 1; + + threadPool.parallelFor(0, (n>>1), [&] (int begin, int end, int numThread) { + for (u_int64_t i=begin; i< end; i++) { + Element t; + Element u; + u_int64_t k=(i/mdiv2)*m; + u_int64_t j=i%mdiv2; + + f.mul(t, root(s, j), a[k+j+mdiv2]); + f.copy(u,a[k+j]); + f.add(a[k+j], t, u); + f.sub(a[k+j+mdiv2], u, t); + } + }); + } +} + +// Inverse of fftDITRevToNat run backwards: decimation in frequency with +// inverse twiddles. Leaves the result scaled by n; the caller folds 1/n +// into its next pointwise pass. +template +void FFT::ifftDIFNatToRev(Element *a, u_int64_t n) { + u_int64_t domainPow = log2(n); + assert(((u_int64_t)1 << domainPow) == n); + + for (u_int32_t s=domainPow; s>=1; s--) { + u_int64_t m = 1 << s; + u_int64_t mdiv2 = m >> 1; + + threadPool.parallelFor(0, (n>>1), [&] (int begin, int end, int numThread) { + for (u_int64_t i=begin; i< end; i++) { + Element t; + Element u; + u_int64_t k=(i/mdiv2)*m; + u_int64_t j=i%mdiv2; + + f.copy(u, a[k+j]); + f.copy(t, a[k+j+mdiv2]); + f.add(a[k+j], u, t); + f.sub(t, u, t); + f.mul(a[k+j+mdiv2], t, rootInv(s, j)); + } + }); + } +} + template void FFT::ifft(Element *a, u_int64_t n ) { fft(a, n); diff --git a/c/fft.hpp b/c/fft.hpp index ab3c63b..21e5d49 100644 --- a/c/fft.hpp +++ b/c/fft.hpp @@ -24,9 +24,28 @@ class FFT { void fft(Element *a, u_int64_t n ); void ifft(Element *a, u_int64_t n ); + // Permutation-free pair: ifftDIFNatToRev takes natural order and leaves + // the (unscaled by 1/n!) inverse transform in bit-reversed order; + // fftDITRevToNat takes bit-reversed order and leaves the forward + // transform in natural order. Chaining them with a pointwise pass in + // between (indexed through BR) round-trips to natural order without any + // bit-reversal permutation of the data. + void ifftDIFNatToRev(Element *a, u_int64_t n); + void fftDITRevToNat(Element *a, u_int64_t n); + u_int32_t log2(u_int64_t n); inline Element &root(u_int32_t domainPow, u_int64_t idx) { return roots[ idx << (s-domainPow)]; } + // Primitive 2^(s+extraPow)-th root of unity — an order finer than the + // table covers (e.g. the omega_2n of a coset shift, without paying for + // a table twice the transform size). Requires s+extraPow within the + // field's 2-adicity. + void higherRootOfUnity(Element &r, u_int32_t extraPow); + inline Element &rootInv(u_int32_t domainPow, u_int64_t idx) { + return roots[ idx == 0 ? 0 : ((((u_int64_t)1 << domainPow) - idx) << (s-domainPow)) ]; + } + inline Element &nInv(u_int32_t domainPow) { return powTwoInv[domainPow]; } + void printVector(Element *a, u_int64_t n ); }; diff --git a/c/msm.cpp b/c/msm.cpp index 204e2cd..3460f1f 100644 --- a/c/msm.cpp +++ b/c/msm.cpp @@ -1,108 +1,584 @@ #include +#include +#include #include "msm.hpp" #include "misc.hpp" template -void MSM::run(typename Curve::Point &r, - typename Curve::PointAffine *_bases, - uint8_t* _scalars, - uint64_t _scalarSize, - uint64_t _n, - uint64_t _nThreads) +void MSM::preparePartition(Partition &p, uint64_t nThreads) { ThreadPool &threadPool = ThreadPool::defaultPool(); - const uint64_t nThreads = threadPool.getThreadCount(); - const uint64_t nPoints = _n; - - scalars = _scalars; - scalarSize = _scalarSize; - #ifdef MSM_BITS_PER_CHUNK - bitsPerChunk = MSM_BITS_PER_CHUNK; + p.bitsPerChunk = MSM_BITS_PER_CHUNK; + p.nSlices = 1; #else - bitsPerChunk = calcBitsPerChunk(nPoints, scalarSize); + calcChunkConfig(p.n, p.nBits, nThreads, p.bitsPerChunk, p.nSlices); #endif - if (nPoints == 0) { - g.copy(r, g.zero()); - return; - } - if (nPoints == 1) { - g.mulByScalar(r, _bases[0], scalars, scalarSize); - return; - } + p.nChunks = calcChunkCount(p.nBits, p.bitsPerChunk); + p.nBuckets = calcBucketCount(p.bitsPerChunk); + p.digits.reset(new int16_t[p.nChunks * p.n]); + p.partials.reset(new typename Curve::Point[p.nSlices * p.nChunks]); + + // Batch-affine pays off only when the bucket array is large (the batch + // stays conflict-free) and densely filled (its two bucket arrays get + // amortized over many additions). + p.batchAffine = (p.bitsPerChunk >= MIN_BATCH_AFFINE_CHUNK_BITS) + && (p.n / p.nSlices >= p.nBuckets); + p.batchSize = std::min(BATCH_SIZE, p.nBuckets/8); - const uint64_t nChunks = calcChunkCount(scalarSize, bitsPerChunk); - const uint64_t nBuckets = calcBucketCount(bitsPerChunk); - const uint64_t matrixSize = nThreads * nBuckets; - const uint64_t nSlices = nChunks*nPoints; + // recode context for getBucketIndex + scalars = p.scalars; + scalarSize = p.scalarSize; + bitsPerChunk = p.bitsPerChunk; - std::unique_ptr bucketMatrix(new typename Curve::Point[matrixSize]); - std::unique_ptr chunks(new typename Curve::Point[nChunks]); - std::unique_ptr slicedScalars(new int32_t[nSlices]); + const uint64_t nChunks = p.nChunks; + const uint64_t nBuckets = p.nBuckets; + const uint64_t nPoints = p.n; + const uint32_t *indices = p.indices; + int16_t *digits = p.digits.get(); - threadPool.parallelFor(0, nPoints, [&] (int begin, int end, int numThread) { + threadPool.parallelFor(0, nPoints, [&, nChunks, nBuckets, nPoints, indices] (int begin, int end, int numThread) { for (int i = begin; i < end; i++) { int carry = 0; + const uint64_t scalarIdx = indices ? indices[i] : (uint64_t)i; - for (int j = 0; j < nChunks; j++) { - int bucketIndex = getBucketIndex(i, j) + carry; + for (uint64_t j = 0; j < nChunks; j++) { + int bucketIndex = getBucketIndex(scalarIdx, j) + carry; - if (bucketIndex >= nBuckets) { + if (bucketIndex >= (int)nBuckets) { bucketIndex -= nBuckets*2; carry = 1; } else { carry = 0; } - slicedScalars[i*nChunks + j] = bucketIndex; + digits[j*nPoints + i] = (int16_t)bucketIndex; } } }); +} - threadPool.parallelFor(0, nChunks, [&] (int begin, int end, int numThread) { +template +void MSM::prepare(typename Curve::PointAffine *_bases, + uint8_t *_scalars, + uint64_t _scalarSize, + uint64_t _n, + uint64_t parallelismShare) +{ + ThreadPool &threadPool = ThreadPool::defaultPool(); - for (int j = begin; j < end; j++) { + const uint64_t nThreads = parallelismShare ? parallelismShare + : threadPool.getThreadCount(); - typename Curve::Point *buckets = &bucketMatrix[numThread*nBuckets]; + partitions.clear(); + partitions.reserve(2); + onesAcc.reset(); + nOnesBlocks = 0; + trivial = false; + prepared = true; - for (int i = 0; i < nBuckets; i++) { - g.copy(buckets[i], g.zero()); - } + if (_n == 0) { + trivial = true; + g.copy(trivialResult, g.zero()); + return; + } + if (_n == 1) { + trivial = true; + g.mulByScalar(trivialResult, _bases[0], _scalars, _scalarSize); + return; + } + + scalars = _scalars; + scalarSize = _scalarSize; - for (int i = 0; i < nPoints; i++) { - const int bucketIndex = slicedScalars[i*nChunks + j]; + if (_scalarSize < 8) { + partitions.emplace_back(); + Partition &p = partitions.back(); + p.indices = NULL; + p.bases = _bases; + p.scalars = _scalars; + p.scalarSize = _scalarSize; + p.n = _n; + p.nBits = _scalarSize*8; + preparePartition(p, nThreads); + return; + } - if (bucketIndex > 0) { - g.add(buckets[bucketIndex-1], buckets[bucketIndex-1], _bases[i]); + const uint64_t nBlocks = std::min(threadPool.getThreadCount()*4, _n); + const uint64_t blockSize = (_n + nBlocks - 1) / nBlocks; - } else if (bucketIndex < 0) { - g.sub(buckets[-bucketIndex-1], buckets[-bucketIndex-1], _bases[i]); + enum ScalarClass : uint8_t { CLS_ZERO = 0, CLS_ONE = 1, CLS_SMALL = 2, CLS_BIG = 3 }; + + std::unique_ptr classes(new uint8_t[_n]); + std::unique_ptr blockCounts(new uint64_t[nBlocks*3]); + std::unique_ptr blockMaxBits(new uint64_t[nBlocks*2]); + + threadPool.parallelFor(0, nBlocks, [&] (int begin, int end, int numThread) { + for (int b = begin; b < end; b++) { + const uint64_t i0 = (uint64_t)b*blockSize; + const uint64_t i1 = std::min(i0 + blockSize, _n); + uint64_t nSmall = 0, nBig = 0, nOnes = 0; + uint64_t maxSmall = 0, maxBig = 0; + + for (uint64_t i = i0; i < i1; i++) { + const uint64_t bits = significantBits(_scalars + i*_scalarSize); + uint8_t cls; + + if (bits == 0) { + cls = CLS_ZERO; + } else if (bits == 1) { + cls = CLS_ONE; + nOnes++; + } else if (bits <= SMALL_SCALAR_BITS) { + cls = CLS_SMALL; + nSmall++; + if (bits > maxSmall) maxSmall = bits; + } else { + cls = CLS_BIG; + nBig++; + if (bits > maxBig) maxBig = bits; } + classes[i] = cls; } + blockCounts[b*3] = nSmall; + blockCounts[b*3+1] = nBig; + blockCounts[b*3+2] = nOnes; + blockMaxBits[b*2] = maxSmall; + blockMaxBits[b*2+1] = maxBig; + } + }); + + uint64_t nSmall = 0, nBig = 0, nOnes = 0; + uint64_t maxSmallBits = 0, maxBigBits = 0; - typename Curve::Point t, tmp; + std::unique_ptr blockOffsets(new uint64_t[nBlocks*2]); - g.copy(t, buckets[nBuckets - 1]); - g.copy(tmp, t); + for (uint64_t b = 0; b < nBlocks; b++) { + blockOffsets[b*2] = nSmall; + blockOffsets[b*2+1] = nBig; + nSmall += blockCounts[b*3]; + nBig += blockCounts[b*3+1]; + nOnes += blockCounts[b*3+2]; + if (blockMaxBits[b*2] > maxSmallBits) maxSmallBits = blockMaxBits[b*2]; + if (blockMaxBits[b*2+1] > maxBigBits) maxBigBits = blockMaxBits[b*2+1]; + } + + const uint64_t overallMaxBits = std::max(maxBigBits, std::max(maxSmallBits, (uint64_t)(nOnes ? 1 : 0))); - for (int i = nBuckets - 2; i >= 0 ; i--) { - g.add(tmp, tmp, buckets[i]); - g.add(t, t, tmp); + // When almost every scalar is full width (e.g. the H MSM, whose scalars + // are uniform field elements) partitioning saves nothing: run the whole + // input in place instead of paying the gather. + if (nBig >= _n - _n/16) { + partitions.emplace_back(); + Partition &p = partitions.back(); + p.indices = NULL; + p.bases = _bases; + p.scalars = _scalars; + p.scalarSize = _scalarSize; + p.n = _n; + p.nBits = overallMaxBits + 2; + preparePartition(p, nThreads); + return; + } + + // All scalars fit in 64 bits: gather only the scalars, bases stay in place. + if (nSmall == _n) { + partitions.emplace_back(); + Partition &p = partitions.back(); + p.ownScalars64.reset(new uint64_t[_n]); + + uint64_t *s64 = p.ownScalars64.get(); + + threadPool.parallelFor(0, _n, [&, s64] (int begin, int end, int numThread) { + for (int i = begin; i < end; i++) { + std::memcpy(&s64[i], _scalars + (uint64_t)i*_scalarSize, sizeof(uint64_t)); } + }); - chunks[j] = t; + p.indices = NULL; + p.bases = _bases; + p.scalars = (uint8_t *)s64; + p.scalarSize = sizeof(uint64_t); + p.n = _n; + p.nBits = maxSmallBits + 2; + preparePartition(p, nThreads); + return; + } + + Partition *small = NULL; + Partition *big = NULL; + + if (nSmall > 0) { + partitions.emplace_back(); + small = &partitions.back(); + small->ownIndices.reset(new uint32_t[nSmall]); + small->indices = small->ownIndices.get(); + small->bases = _bases; + small->scalars = _scalars; + small->scalarSize = _scalarSize; + small->n = nSmall; + small->nBits = maxSmallBits + 2; + } + if (nBig > 0) { + partitions.emplace_back(); + big = &partitions.back(); + big->ownIndices.reset(new uint32_t[nBig]); + big->indices = big->ownIndices.get(); + big->bases = _bases; + big->scalars = _scalars; + big->scalarSize = _scalarSize; + big->n = nBig; + big->nBits = maxBigBits + 2; + } + + nOnesBlocks = nBlocks; + onesAcc.reset(new typename Curve::Point[nBlocks]); + + typename Curve::Point *ones = onesAcc.get(); + + threadPool.parallelFor(0, nBlocks, [&, ones] (int begin, int end, int numThread) { + for (int b = begin; b < end; b++) { + const uint64_t i0 = (uint64_t)b*blockSize; + const uint64_t i1 = std::min(i0 + blockSize, _n); + uint64_t smallCur = blockOffsets[b*2]; + uint64_t bigCur = blockOffsets[b*2+1]; + + g.copy(ones[b], g.zero()); + + for (uint64_t i = i0; i < i1; i++) { + switch (classes[i]) { + case CLS_ONE: + g.add(ones[b], ones[b], _bases[i]); + break; + case CLS_SMALL: + small->ownIndices[smallCur++] = (uint32_t)i; + break; + case CLS_BIG: + big->ownIndices[bigCur++] = (uint32_t)i; + break; + default: + break; + } + } } }); - g.copy(r, chunks[nChunks - 1]); + if (small) preparePartition(*small, nThreads); + if (big) preparePartition(*big, nThreads); +} + +template +uint64_t MSM::arenaBytesPerThread() const +{ + uint64_t m = 0; + + for (const Partition &p : partitions) { + const uint64_t bytes = partitionArenaBytes(p); + if (bytes > m) m = bytes; + } + return (m + 63) & ~(uint64_t)63; +} + +template +void MSM::fillChunkXYZZ(Partition &p, uint64_t j, + uint64_t i0, uint64_t i1, + uint64_t sliceIdx, uint8_t *taskArena) +{ + typename Curve::Point *buckets = (typename Curve::Point *)taskArena; + const int16_t *digits = &p.digits[j*p.n]; + typename Curve::PointAffine *bases = p.bases; + const uint32_t *indices = p.indices; + const uint64_t nBuckets = p.nBuckets; + + for (uint64_t i = 0; i < nBuckets; i++) { + g.copy(buckets[i], g.zero()); + } + + for (uint64_t i = i0; i < i1; i++) { + const int32_t bucketIndex = digits[i]; + typename Curve::PointAffine &base = bases[indices ? indices[i] : i]; - for (int j = nChunks - 2; j >= 0; j--) { - for (int i = 0; i < bitsPerChunk; i++) { - g.dbl(r, r); + if (bucketIndex > 0) { + g.add(buckets[bucketIndex-1], buckets[bucketIndex-1], base); + + } else if (bucketIndex < 0) { + g.sub(buckets[-bucketIndex-1], buckets[-bucketIndex-1], base); } - g.add(r, r, chunks[j]); } + + typename Curve::Point t, tmp; + + g.copy(t, buckets[nBuckets - 1]); + g.copy(tmp, t); + + for (int64_t i = nBuckets - 2; i >= 0 ; i--) { + g.add(tmp, tmp, buckets[i]); + g.add(t, t, tmp); + } + + p.partials[sliceIdx*p.nChunks + j] = t; +} + +template +void MSM::fillChunkBatchAffine(Partition &p, uint64_t j, + uint64_t i0, uint64_t i1, + uint64_t sliceIdx, uint8_t *taskArena) +{ + typedef typename Curve::PointAffine PointAffine; + typedef typename Curve::Point Point; + typedef typename BaseField::Element Element; + + const uint64_t nBuckets = p.nBuckets; + const uint64_t batchSize = p.batchSize; + const int16_t *digits = &p.digits[j*p.n]; + PointAffine *bases = p.bases; + const uint32_t *indices = p.indices; + BaseField &F = g.F; + + uint8_t *cur = taskArena; + PointAffine *buckets = (PointAffine *)cur; cur += nBuckets*sizeof(PointAffine); + Point *shadow = (Point *)cur; cur += nBuckets*sizeof(Point); + PointAffine *batchP = (PointAffine *)cur; cur += batchSize*sizeof(PointAffine); + Element *dx = (Element *)cur; cur += batchSize*sizeof(Element); + Element *prod = (Element *)cur; cur += batchSize*sizeof(Element); + uint32_t *batchB = (uint32_t *)cur; cur += batchSize*sizeof(uint32_t); + uint8_t *inBatch = cur; // nBuckets bytes + + // all-zero bytes encode infinity in both representations + std::memset(buckets, 0, nBuckets*sizeof(PointAffine)); + std::memset(shadow, 0, nBuckets*sizeof(Point)); + std::memset(inBatch, 0, nBuckets); + + uint64_t count = 0; + + // Execute the pending independent affine additions, amortizing one + // inversion over the whole batch (Montgomery's trick). + auto executeBatch = [&] () { + if (count == 0) return; + + for (uint64_t k = 0; k < count; k++) { + F.sub(dx[k], batchP[k].x, buckets[batchB[k]].x); + + if (k == 0) { + F.copy(prod[0], dx[0]); + } else { + F.mul(prod[k], prod[k-1], dx[k]); + } + } + + Element invAll, invK, lambda, t1, x3; + + F.inv(invAll, prod[count-1]); + + for (int64_t k = count - 1; k >= 0; k--) { + PointAffine &B = buckets[batchB[k]]; + + if (k > 0) { + F.mul(invK, invAll, prod[k-1]); + F.mul(invAll, invAll, dx[k]); + } else { + F.copy(invK, invAll); + } + + // chord addition: B = B + P + F.sub(t1, batchP[k].y, B.y); + F.mul(lambda, t1, invK); + + F.square(x3, lambda); + F.sub(x3, x3, B.x); + F.sub(x3, x3, batchP[k].x); + + F.sub(t1, B.x, x3); + F.mul(t1, t1, lambda); + F.sub(B.y, t1, B.y); + F.copy(B.x, x3); + + inBatch[batchB[k]] = 0; + } + count = 0; + }; + + for (uint64_t i = i0; i < i1; i++) { + const int32_t d = digits[i]; + + if (d == 0) continue; + + PointAffine &base = bases[indices ? indices[i] : i]; + + if (g.isZero(base)) continue; + + const uint32_t b = (uint32_t)(d > 0 ? d : -d) - 1; + + PointAffine P; + F.copy(P.x, base.x); + if (d > 0) { + F.copy(P.y, base.y); + } else { + F.neg(P.y, base.y); + } + + if (inBatch[b]) { + // the bucket has a pending addition: divert to its shadow + g.add(shadow[b], shadow[b], P); + continue; + } + if (F.isZero(buckets[b].x) && F.isZero(buckets[b].y)) { + buckets[b] = P; + continue; + } + if (F.eq(buckets[b].x, P.x)) { + if (F.eq(buckets[b].y, P.y)) { + // doubling: fold 2P into the shadow bucket + Point t2; + g.dbl(t2, P); + g.add(shadow[b], shadow[b], t2); + } + // else P == -bucket: they cancel + std::memset(&buckets[b], 0, sizeof(PointAffine)); + continue; + } + + batchB[count] = b; + batchP[count] = P; + inBatch[b] = 1; + count++; + + if (count == batchSize) executeBatch(); + } + executeBatch(); + + typename Curve::Point t, tmp; + + g.copy(t, g.zero()); + g.copy(tmp, g.zero()); + + for (int64_t b = nBuckets - 1; b >= 0; b--) { + if (!(F.isZero(buckets[b].x) && F.isZero(buckets[b].y))) { + g.add(tmp, tmp, buckets[b]); + } + if (!g.isZero(shadow[b])) { + g.add(tmp, tmp, shadow[b]); + } + g.add(t, t, tmp); + } + + p.partials[sliceIdx*p.nChunks + j] = t; +} + +template +void MSM::collectTasks(std::vector &tasks, + uint8_t *bucketArena, + uint64_t bytesPerThread) +{ + for (Partition &part : partitions) { + Partition *p = ∂ + + for (uint64_t s = 0; s < p->nSlices; s++) { + const uint64_t i0 = p->n * s / p->nSlices; + const uint64_t i1 = p->n * (s+1) / p->nSlices; + + for (uint64_t j = 0; j < p->nChunks; j++) { + tasks.push_back([this, p, s, j, i0, i1, bucketArena, bytesPerThread] (uint64_t threadId) { + uint8_t *taskArena = bucketArena + threadId*bytesPerThread; + + if (p->batchAffine) { + fillChunkBatchAffine(*p, j, i0, i1, s, taskArena); + } else { + fillChunkXYZZ(*p, j, i0, i1, s, taskArena); + } + }); + } + } + } +} + +template +void MSM::reducePartition(Partition &p, typename Curve::Point &r) +{ + typename Curve::Point chunkSum; + + for (int64_t j = p.nChunks - 1; j >= 0; j--) { + g.copy(chunkSum, p.partials[j]); + for (uint64_t s = 1; s < p.nSlices; s++) { + g.add(chunkSum, chunkSum, p.partials[s*p.nChunks + j]); + } + + if (j == (int64_t)p.nChunks - 1) { + g.copy(r, chunkSum); + } else { + g.add(r, r, chunkSum); + } + + if (j > 0) { + for (uint64_t b = 0; b < p.bitsPerChunk; b++) { + g.dbl(r, r); + } + } + } +} + +template +void MSM::finish(typename Curve::Point &r) +{ + if (trivial) { + g.copy(r, trivialResult); + prepared = false; + return; + } + + typename Curve::Point acc, part; + + g.copy(acc, g.zero()); + + for (Partition &p : partitions) { + reducePartition(p, part); + g.add(acc, acc, part); + } + + for (uint64_t b = 0; b < nOnesBlocks; b++) { + g.add(acc, acc, onesAcc[b]); + } + + g.copy(r, acc); + + partitions.clear(); + onesAcc.reset(); + nOnesBlocks = 0; + prepared = false; +} + +template +void MSM::run(typename Curve::Point &r, + typename Curve::PointAffine *_bases, + uint8_t* _scalars, + uint64_t _scalarSize, + uint64_t _n, + uint64_t _nThreads) +{ + ThreadPool &threadPool = ThreadPool::defaultPool(); + + prepare(_bases, _scalars, _scalarSize, _n); + + if (!trivial) { + const uint64_t nThreads = threadPool.getThreadCount(); + const uint64_t bytesPerThread = arenaBytesPerThread(); + + std::unique_ptr arena(new uint8_t[nThreads * bytesPerThread]); + + std::vector tasks; + collectTasks(tasks, arena.get(), bytesPerThread); + + if (!tasks.empty()) { + threadPool.parallelFor(0, tasks.size(), [&] (int begin, int end, int numThread) { + for (int t = begin; t < end; t++) { + tasks[t]((uint64_t)numThread); + } + }); + } + } + + finish(r); } diff --git a/c/msm.hpp b/c/msm.hpp index bdfdb69..efb2d71 100644 --- a/c/msm.hpp +++ b/c/msm.hpp @@ -2,48 +2,152 @@ #define MSM_HPP #include - +#include +#include +#include + +// Pippenger bucket-method MSM with scalar-size partitioning and a +// task-based execution model. +// +// Usage (single MSM): run() — prepare, execute and reduce in one call. +// +// Usage (batched, e.g. the prover's A/B1/B2/C phase): +// msm.prepare(bases, scalars, size, n, parallelismShare); +// msm.collectTasks(tasks, arena); // arena: nThreads*maxBuckets() Points +// ... run all MSMs' tasks in one parallel region ... +// msm.finish(r); +// Tasks from several MSM instances (even over different curves) can be +// mixed in one region; each task gets the executing thread id and uses +// that thread's row of its own curve's bucket arena. template class MSM { +public: + typedef std::function Task; + +private: const uint64_t MIN_CHUNK_SIZE_BITS = 3; const uint64_t MAX_CHUNK_SIZE_BITS = 16; + // (enum members: in-class integral constants that are never ODR-used, + // so no out-of-class definitions are needed) + enum : uint64_t { + // Scalars of at most this many significant bits go to the small + // partition, which runs the bucket method over a single 64-bit word + // and so pays ~4 windows instead of ~16. Scalars 0 and 1 are cheaper + // still: they need no scalar multiplication at all. + SMALL_SCALAR_BITS = 64, + + // Don't point-split below this many points per slice: the per-slice + // running-sum cost (2^c bucket additions) would dominate. + MIN_POINTS_PER_SLICE = 4096, + + // Batch-affine accumulation: buckets live in affine coordinates and + // additions are executed in batches sharing one field inversion + // (~5M+1S per addition instead of 9M+2S for an XYZZ mixed add). + // Additions that conflict with the pending batch, doublings and + // cancellations go to an XYZZ shadow bucket instead. Used only when + // the bucket array is large and densely filled enough. + BATCH_SIZE = 512, + MIN_BATCH_AFFINE_CHUNK_BITS = 10 + }; + + // One scalar-size class of the input, ready for bucket accumulation. + struct Partition { + typename Curve::PointAffine *bases; // points (caller's or gathered) + uint8_t *scalars; // scalars (caller's or gathered) + uint64_t scalarSize; + uint64_t n; + uint64_t nBits; // significant bits + carry headroom + uint64_t bitsPerChunk; + uint64_t nChunks; + uint64_t nBuckets; + uint64_t nSlices; // point-split factor + bool batchAffine; // bucket accumulation strategy + uint64_t batchSize; + std::unique_ptr digits; // chunk-major [nChunks][n]; signed digits + // fit int16 exactly: |digit| <= 2^(c-1) <= 2^15 for c <= 16 + std::unique_ptr partials; // [nSlices][nChunks] + + // When set, the partition addresses the caller's bases/scalars + // through this ascending index list instead of gathered copies + // (4 bytes per point instead of a point + scalar copy). + const uint32_t *indices; + + // backing storage for indexed/gathered classes + std::unique_ptr ownIndices; + std::unique_ptr ownScalars64; + }; + Curve &g; + + // significantBits()/getBucketIndex() context for the classification and + // recode passes; set before each pass. uint8_t *scalars; uint64_t scalarSize; uint64_t bitsPerChunk; -private: - uint64_t calcAddsCount(uint64_t nPoints, uint64_t scalarSize, uint64_t bitsPerChunk) const { - return calcChunkCount(scalarSize, bitsPerChunk) - * (nPoints + ((uint64_t)1 << bitsPerChunk) + bitsPerChunk + 1); - } + std::vector partitions; + std::unique_ptr onesAcc; // per-block partial sums of 1-scalar points + uint64_t nOnesBlocks; + bool prepared; - uint64_t calcBitsPerChunk(uint64_t n, uint64_t scalarSize) const { - uint64_t bitsPerChunk = MIN_CHUNK_SIZE_BITS; - uint64_t minAdds = calcAddsCount(n, scalarSize, bitsPerChunk); + // Set when prepare() resolved the whole MSM without bucket work + // (n==0, n==1, or every scalar in {0,1}). + bool trivial; + typename Curve::Point trivialResult; - for (uint64_t k = MIN_CHUNK_SIZE_BITS + 1; k <= MAX_CHUNK_SIZE_BITS; k++) { - const uint64_t curAdds = calcAddsCount(n, scalarSize, k); +private: + uint64_t calcChunkCount(uint64_t nBits, uint64_t bitsPerChunk) const { + return ((nBits - 1) / bitsPerChunk) + 1; + } - if (curAdds < minAdds) { - minAdds = curAdds; - bitsPerChunk = k; - } - } - return bitsPerChunk; + uint64_t calcBucketCount(uint64_t bitsPerChunk) const { + return ((uint64_t)1 << (bitsPerChunk-1)); } - uint64_t calcChunkCount(uint64_t scalarSize, uint64_t bitsPerChunk) const { - return ((scalarSize * 8 - 1 ) / bitsPerChunk) + 1; + // Estimated wall-clock cost in point additions of one partition executed + // as nSlices*nChunks tasks on nThreads threads. Tasks run in waves; the + // second term charges 1/8 of the total work so that among near-equal + // wall costs the one burning fewer total additions wins. + uint64_t calcCost(uint64_t n, uint64_t nBits, uint64_t bitsPerChunk, + uint64_t nSlices, uint64_t nThreads) const { + const uint64_t sliceCost = n/nSlices + ((uint64_t)1 << bitsPerChunk) + bitsPerChunk + 1; + const uint64_t nTasks = nSlices * calcChunkCount(nBits, bitsPerChunk); + const uint64_t waves = (nTasks + nThreads - 1) / nThreads; + + return waves*sliceCost + nTasks*sliceCost/(8*nThreads); } - uint64_t calcBucketCount(uint64_t bitsPerChunk) const { - return ((uint64_t)1 << (bitsPerChunk-1)); + // Pick window size and point-split factor minimizing estimated wall cost. + void calcChunkConfig(uint64_t n, uint64_t nBits, uint64_t nThreads, + uint64_t &bestC, uint64_t &bestSlices) const { + const uint64_t maxSlices = std::max(1, std::min( + 2*nThreads, n / MIN_POINTS_PER_SLICE)); + + bestC = MIN_CHUNK_SIZE_BITS; + bestSlices = 1; + uint64_t minCost = calcCost(n, nBits, bestC, 1, nThreads); + + for (uint64_t k = MIN_CHUNK_SIZE_BITS; k <= MAX_CHUNK_SIZE_BITS; k++) { + for (uint64_t s = 1; s <= maxSlices; s *= 2) { + const uint64_t curCost = calcCost(n, nBits, k, s, nThreads); + + if (curCost < minCost) { + minCost = curCost; + bestC = k; + bestSlices = s; + } + } + } } uint64_t getBucketIndex(uint64_t scalarIdx, uint64_t chunkIdx) const { uint64_t bitStart = chunkIdx*bitsPerChunk; + + // Chunks past the scalar bytes exist only to absorb the signed-digit + // carry; their digit is zero. + if (bitStart >= scalarSize*8) return 0; + uint64_t byteStart = bitStart/8; uint64_t efectiveBitsPerChunk = bitsPerChunk; @@ -59,9 +163,71 @@ class MSM { return uint64_t(v); } -public: - MSM(Curve &_g): g(_g) {} + // Number of significant bits of a little-endian scalar; 0 for a zero scalar. + uint64_t significantBits(const uint8_t *scalar) const { + for (int64_t k = (int64_t)scalarSize - 1; k >= 0; k--) { + if (scalar[k]) { + return (uint64_t)k*8 + (32 - __builtin_clz((uint32_t)scalar[k])); + } + } + return 0; + } + // Recode a partition's scalars into signed digits and size its partials. + void preparePartition(Partition &p, uint64_t nThreads); + + // Reduce one partition's task partials into a single point. + void reducePartition(Partition &p, typename Curve::Point &r); + + // Bucket accumulation + running sum for points [i0,i1) of chunk j, + // writing the result into the (slice, chunk) partial. + void fillChunkXYZZ(Partition &p, uint64_t j, uint64_t i0, uint64_t i1, + uint64_t sliceIdx, uint8_t *taskArena); + void fillChunkBatchAffine(Partition &p, uint64_t j, uint64_t i0, uint64_t i1, + uint64_t sliceIdx, uint8_t *taskArena); + + uint64_t partitionArenaBytes(const Partition &p) const { + if (!p.batchAffine) { + return p.nBuckets * sizeof(typename Curve::Point); + } + return p.nBuckets * (sizeof(typename Curve::PointAffine) + + sizeof(typename Curve::Point) + 1) + + p.batchSize * (sizeof(typename Curve::PointAffine) + + 2*sizeof(typename BaseField::Element) + + sizeof(uint32_t)) + + 64; // alignment slack + } + +public: + MSM(Curve &_g): g(_g), prepared(false), trivial(false) {} + + // Classify scalars by size, gather the classes and recode digits. + // parallelismShare: number of threads this MSM should assume it has for + // itself when sizing windows/slices — pass the pool size when the MSM + // runs alone, or roughly poolSize/nMSMs when batched with others. + void prepare(typename Curve::PointAffine *bases, + uint8_t *scalars, + uint64_t scalarSize, + uint64_t n, + uint64_t parallelismShare = 0); + + // Largest per-thread scratch any of this MSM's tasks needs; the caller + // provides an arena of nThreads*arenaBytesPerThread() bytes to + // collectTasks() (8-byte aligned, e.g. from new uint8_t[]). + uint64_t arenaBytesPerThread() const; + + // Append one task per (partition, slice, chunk). Tasks only touch their + // own partials and bucketArena[threadId*bytesPerThread..]. When several + // MSMs share an arena, bytesPerThread is the max of their + // arenaBytesPerThread(). + void collectTasks(std::vector &tasks, + uint8_t *bucketArena, + uint64_t bytesPerThread); + + // Reduce all partials into the final result. Call after every task ran. + void finish(typename Curve::Point &r); + + // Single-MSM convenience: prepare + run own tasks + finish. void run(typename Curve::Point &r, typename Curve::PointAffine *_bases, uint8_t* _scalars,