Rewrite logic to process more than 127 bytes at a time (#14)

Increase the 'carry' size, decrease the num of queues and add ability to
push the same piece of memory through the layers. While the code could
be made even smarter, pushing this version is a good starting point.
This change gives a 4x speedup.

Author: olcbean

commp.go

@@ -26,10 +26,10 @@ type Calc struct {
 	mu sync.Mutex
 }
 type state struct {
-	bytesConsumed uint64
+	quadsEnqueued uint64
 	layerQueues   [MaxLayers + 2]chan []byte // one extra layer for the initial leaves, one more for the dummy never-to-use channel
 	resultCommP   chan []byte
-	carry         []byte
+	buffer        []byte
 }
 
 var _ hash.Hash = &Calc{} // make sure we are hash.Hash compliant
@@ -50,42 +50,49 @@ const MaxPiecePayload = MaxPieceSize / 128 * 127
 // at least this amount of bytes.
 const MinPiecePayload = uint64(65)
 
+const (
+	commpDigestSize = sha256simd.Size
+	quadPayload     = 127
+	bufferSize      = 256 * quadPayload // FIXME: tune better, chosen by rough experiment
+)
+
 var (
-	layerQueueDepth   = 256 // SANCHECK: too much? too little? can't think this through right now...
+	layerQueueDepth   = 32 // FIXME: tune better, chosen by rough experiment
 	shaPool           = sync.Pool{New: func() interface{} { return sha256simd.New() }}
 	stackedNulPadding [MaxLayers][]byte
 )
 
 // initialize the nul padding stack (cheap to do upfront, just MaxLayers loops)
 func init() {
 	h := shaPool.Get().(hash.Hash)
-	defer shaPool.Put(h)
 
-	stackedNulPadding[0] = make([]byte, 32)
+	stackedNulPadding[0] = make([]byte, commpDigestSize)
 	for i := uint(1); i < MaxLayers; i++ {
 		h.Reset()
-		h.Write(stackedNulPadding[i-1]) // yes, got to
-		h.Write(stackedNulPadding[i-1]) // do it twice
-		stackedNulPadding[i] = h.Sum(make([]byte, 0, 32))
+		h.Write(stackedNulPadding[i-1]) // yes, got to...
+		h.Write(stackedNulPadding[i-1]) // ...do it twice
+		stackedNulPadding[i] = h.Sum(make([]byte, 0, commpDigestSize))
 		stackedNulPadding[i][31] &= 0x3F
 	}
+
+	shaPool.Put(h)
 }
 
 // BlockSize is the amount of bytes consumed by the commP algorithm in one go.
 // Write()ing data in multiples of BlockSize would obviate the need to maintain
 // an internal carry buffer. The BlockSize of this module is 127 bytes.
-func (cp *Calc) BlockSize() int { return 127 }
+func (cp *Calc) BlockSize() int { return quadPayload }
 
 // Size is the amount of bytes returned on Sum()/Digest(), which is 32 bytes
 // for this module.
-func (cp *Calc) Size() int { return 32 }
+func (cp *Calc) Size() int { return commpDigestSize }
 
 // Reset re-initializes the accumulator object, clearing its state and
 // terminating all background goroutines. It is safe to Reset() an accumulator
 // in any state.
 func (cp *Calc) Reset() {
 	cp.mu.Lock()
-	if cp.bytesConsumed != 0 {
+	if cp.buffer != nil {
 		// we are resetting without digesting: close everything out to terminate
 		// the layer workers
 		close(cp.layerQueues[0])
@@ -123,28 +130,33 @@ func (cp *Calc) Digest() (commP []byte, paddedPieceSize uint64, err error) {
 		cp.mu.Unlock()
 	}()
 
-	if cp.bytesConsumed < MinPiecePayload {
+	if processed := cp.quadsEnqueued*quadPayload + uint64(len(cp.buffer)); processed < MinPiecePayload {
 		err = xerrors.Errorf(
 			"insufficient state accumulated: commP is not defined for inputs shorter than %d bytes, but only %d processed so far",
-			MinPiecePayload, cp.bytesConsumed,
+			MinPiecePayload, processed,
 		)
 		return
 	}
 
 	// If any, flush remaining bytes padded up with zeroes
-	if len(cp.carry) > 0 {
-		if len(cp.carry) < 127 {
-			cp.carry = append(cp.carry, make([]byte, 127-len(cp.carry))...)
+	if len(cp.buffer) > 0 {
+		if mod := len(cp.buffer) % quadPayload; mod != 0 {
+			cp.buffer = append(cp.buffer, make([]byte, quadPayload-mod)...)
+		}
+		for len(cp.buffer) > 0 {
+			// FIXME: there is a smarter way to do this instead of 127-at-a-time,
+			// but that's for another PR
+			cp.digestQuads(cp.buffer[:127])
+			cp.buffer = cp.buffer[127:]
 		}
-		cp.digestLeading127Bytes(cp.carry)
 	}
 
 	// This is how we signal to the bottom of the stack that we are done
 	// which in turn collapses the rest all the way to resultCommP
 	close(cp.layerQueues[0])
 
+	paddedPieceSize = cp.quadsEnqueued * 128
 	// hacky round-up-to-next-pow2
-	paddedPieceSize = ((cp.bytesConsumed + 126) / 127 * 128) // why is 6 afraid of 7...?
 	if bits.OnesCount64(paddedPieceSize) != 1 {
 		paddedPieceSize = 1 << uint(64-bits.LeadingZeros64(paddedPieceSize))
 	}
@@ -161,115 +173,110 @@ func (cp *Calc) Digest() (commP []byte, paddedPieceSize uint64, err error) {
 // amount of bytes is about to go over the maximum currently supported by
 // Filecoin.
 func (cp *Calc) Write(input []byte) (int, error) {
-	inputSize := len(input)
-	if inputSize == 0 {
+	if len(input) == 0 {
 		return 0, nil
 	}
 
 	cp.mu.Lock()
 	defer cp.mu.Unlock()
 
-	if cp.bytesConsumed+uint64(inputSize) > MaxPiecePayload {
+	if MaxPiecePayload <
+		(cp.quadsEnqueued*quadPayload)+
+			uint64(len(input)) {
 		return 0, xerrors.Errorf(
-			"writing %d bytes to the accumulator would overflow the maximum supported unpadded piece size %d",
+			"writing additional %d bytes to the accumulator would overflow the maximum supported unpadded piece size %d",
 			len(input), MaxPiecePayload,
 		)
 	}
 
 	// just starting: initialize internal state, start first background layer-goroutine
-	if cp.bytesConsumed == 0 {
-		cp.carry = make([]byte, 0, 127)
+	if cp.buffer == nil {
+		cp.buffer = make([]byte, 0, bufferSize)
 		cp.resultCommP = make(chan []byte, 1)
 		cp.layerQueues[0] = make(chan []byte, layerQueueDepth)
 		cp.addLayer(0)
 	}
 
-	cp.bytesConsumed += uint64(inputSize)
-
-	carrySize := len(cp.carry)
-	if carrySize > 0 {
+	// short Write() - just buffer it
+	if len(cp.buffer)+len(input) < bufferSize {
+		cp.buffer = append(cp.buffer, input...)
+		return len(input), nil
+	}
 
-		// super short Write - just carry it
-		if carrySize+inputSize < 127 {
-			cp.carry = append(cp.carry, input...)
-			return inputSize, nil
-		}
+	totalInputBytes := len(input)
 
-		cp.carry = append(cp.carry, input[:127-carrySize]...)
-		input = input[127-carrySize:]
+	if toSplice := bufferSize - len(cp.buffer); toSplice < bufferSize {
+		cp.buffer = append(cp.buffer, input[:toSplice]...)
+		input = input[toSplice:]
 
-		cp.digestLeading127Bytes(cp.carry)
-		cp.carry = cp.carry[:0]
+		cp.digestQuads(cp.buffer)
+		cp.buffer = cp.buffer[:0]
 	}
 
-	for len(input) >= 127 {
-		cp.digestLeading127Bytes(input)
-		input = input[127:]
+	for len(input) >= bufferSize {
+		cp.digestQuads(input[:bufferSize])
+		input = input[bufferSize:]
 	}
 
 	if len(input) > 0 {
-		cp.carry = cp.carry[:len(input)]
-		copy(cp.carry, input)
+		cp.buffer = append(cp.buffer, input...)
 	}
 
-	return inputSize, nil
+	return totalInputBytes, nil
 }
 
-func (cp *Calc) digestLeading127Bytes(input []byte) {
-
-	// Holds this round's shifts of the original 127 bytes plus the 6 bit overflow
-	// at the end of the expansion cycle. We *do not* reuse this array: it is
-	// being fed piece-wise to hash254Into which in turn reuses it for the result
-	var expander [128]byte
-
-	// Cycle over four(4) 31-byte groups, leaving 1 byte in between:
-	// 31 + 1 + 31 + 1 + 31 + 1 + 31 = 127
-
-	// First 31 bytes + 6 bits are taken as-is (trimmed later)
-	// Note that copying them into the expansion buffer is mandatory:
-	// we will be feeding it to the workers which reuse the bottom half
-	// of the chunk for the result
-	copy(expander[:], input[:32])
-
-	// first 2-bit "shim" forced into the otherwise identical bitstream
-	expander[31] &= 0x3F
-
-	// simplify pointer math
-	inputPlus1, expanderPlus1 := input[1:], expander[1:]
-
-	//  In: {{ C[7] C[6] }} X[7] X[6] X[5] X[4] X[3] X[2] X[1] X[0] Y[7] Y[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] Z[7] Z[6] Z[5]...
-	// Out:                 X[5] X[4] X[3] X[2] X[1] X[0] C[7] C[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] X[7] X[6] Z[5] Z[4] Z[3]...
-	for i := 31; i < 63; i++ {
-		expanderPlus1[i] = inputPlus1[i]<<2 | input[i]>>6
-	}
+// always called with power-of-2 amount of quads
+func (cp *Calc) digestQuads(inSlab []byte) {
+
+	quadsCount := len(inSlab) / 127
+	cp.quadsEnqueued += uint64(quadsCount)
+	outSlab := make([]byte, quadsCount*128)
+
+	for j := 0; j < quadsCount; j++ {
+		// Cycle over four(4) 31-byte groups, leaving 1 byte in between:
+		// 31 + 1 + 31 + 1 + 31 + 1 + 31 = 127
+		input := inSlab[j*127 : (j+1)*127]
+		expander := outSlab[j*128 : (j+1)*128]
+		inputPlus1, expanderPlus1 := input[1:], expander[1:]
+
+		// First 31 bytes + 6 bits are taken as-is (trimmed later)
+		// Note that copying them into the expansion buffer is mandatory:
+		// we will be feeding it to the workers which reuse the bottom half
+		// of the chunk for the result
+		copy(expander[:], input[:32])
+
+		// first 2-bit "shim" forced into the otherwise identical bitstream
+		expander[31] &= 0x3F
+
+		//  In: {{ C[7] C[6] }} X[7] X[6] X[5] X[4] X[3] X[2] X[1] X[0] Y[7] Y[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] Z[7] Z[6] Z[5]...
+		// Out:                 X[5] X[4] X[3] X[2] X[1] X[0] C[7] C[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] X[7] X[6] Z[5] Z[4] Z[3]...
+		for i := 31; i < 63; i++ {
+			expanderPlus1[i] = inputPlus1[i]<<2 | input[i]>>6
+		}
 
-	// next 2-bit shim
-	expander[63] &= 0x3F
+		// next 2-bit shim
+		expander[63] &= 0x3F
 
-	// ready to dispatch first half
-	cp.layerQueues[0] <- expander[0:32]
-	cp.layerQueues[0] <- expander[32:64]
+		//  In: {{ C[7] C[6] C[5] C[4] }} X[7] X[6] X[5] X[4] X[3] X[2] X[1] X[0] Y[7] Y[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] Z[7] Z[6] Z[5]...
+		// Out:                           X[3] X[2] X[1] X[0] C[7] C[6] C[5] C[4] Y[3] Y[2] Y[1] Y[0] X[7] X[6] X[5] X[4] Z[3] Z[2] Z[1]...
+		for i := 63; i < 95; i++ {
+			expanderPlus1[i] = inputPlus1[i]<<4 | input[i]>>4
+		}
 
-	//  In: {{ C[7] C[6] C[5] C[4] }} X[7] X[6] X[5] X[4] X[3] X[2] X[1] X[0] Y[7] Y[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] Z[7] Z[6] Z[5]...
-	// Out:                           X[3] X[2] X[1] X[0] C[7] C[6] C[5] C[4] Y[3] Y[2] Y[1] Y[0] X[7] X[6] X[5] X[4] Z[3] Z[2] Z[1]...
-	for i := 63; i < 95; i++ {
-		expanderPlus1[i] = inputPlus1[i]<<4 | input[i]>>4
-	}
+		// next 2-bit shim
+		expander[95] &= 0x3F
 
-	// next 2-bit shim
-	expander[95] &= 0x3F
+		//  In: {{ C[7] C[6] C[5] C[4] C[3] C[2] }} X[7] X[6] X[5] X[4] X[3] X[2] X[1] X[0] Y[7] Y[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] Z[7] Z[6] Z[5]...
+		// Out:                                     X[1] X[0] C[7] C[6] C[5] C[4] C[3] C[2] Y[1] Y[0] X[7] X[6] X[5] X[4] X[3] X[2] Z[1] Z[0] Y[7]...
+		for i := 95; i < 126; i++ {
+			expanderPlus1[i] = inputPlus1[i]<<6 | input[i]>>2
+		}
 
-	//  In: {{ C[7] C[6] C[5] C[4] C[3] C[2] }} X[7] X[6] X[5] X[4] X[3] X[2] X[1] X[0] Y[7] Y[6] Y[5] Y[4] Y[3] Y[2] Y[1] Y[0] Z[7] Z[6] Z[5]...
-	// Out:                                     X[1] X[0] C[7] C[6] C[5] C[4] C[3] C[2] Y[1] Y[0] X[7] X[6] X[5] X[4] X[3] X[2] Z[1] Z[0] Y[7]...
-	for i := 95; i < 126; i++ {
-		expanderPlus1[i] = inputPlus1[i]<<6 | input[i]>>2
+		// the final 6 bit remainder is exactly the value of the last expanded byte
+		expander[127] = input[126] >> 2
 	}
-	// the final 6 bit remainder is exactly the value of the last expanded byte
-	expander[127] = input[126] >> 2
 
-	// and dispatch remainder
-	cp.layerQueues[0] <- expander[64:96]
-	cp.layerQueues[0] <- expander[96:128]
+	cp.layerQueues[0] <- outSlab
 }
 
 func (cp *Calc) addLayer(myIdx uint) {
@@ -280,60 +287,71 @@ func (cp *Calc) addLayer(myIdx uint) {
 	cp.layerQueues[myIdx+1] = make(chan []byte, layerQueueDepth)
 
 	go func() {
-		var chunkHold []byte
+		var twinHold []byte
 
 		for {
-
-			chunk, queueIsOpen := <-cp.layerQueues[myIdx]
+			slab, queueIsOpen := <-cp.layerQueues[myIdx]
 
 			// the dream is collapsing
 			if !queueIsOpen {
+				defer func() { twinHold = nil }()
 
 				// I am last
 				if myIdx == MaxLayers || cp.layerQueues[myIdx+2] == nil {
-					cp.resultCommP <- chunkHold
+					cp.resultCommP <- append(make([]byte, 0, 32), twinHold[0:32]...)
 					return
 				}
 
-				if chunkHold != nil {
-					cp.hash254Into(
-						cp.layerQueues[myIdx+1],
-						chunkHold,
-						stackedNulPadding[myIdx],
-					)
+				if twinHold != nil {
+					copy(twinHold[32:64], stackedNulPadding[myIdx])
+					cp.hashSlab254(0, twinHold[0:64])
+					cp.layerQueues[myIdx+1] <- twinHold[0:64:64]
 				}
 
 				// signal the next in line that they are done too
 				close(cp.layerQueues[myIdx+1])
 				return
 			}
 
-			if chunkHold == nil {
-				chunkHold = chunk
-			} else {
-
-				// We are last right now
-				// n.b. we will not blow out of the preallocated layerQueues array,
-				// as we disallow Write()s above a certain threshold
-				if cp.layerQueues[myIdx+2] == nil {
-					cp.addLayer(myIdx + 1)
-				}
+			var pushedWork bool
+
+			switch {
+			case len(slab) > 1<<(5+myIdx):
+				cp.hashSlab254(myIdx, slab)
+				cp.layerQueues[myIdx+1] <- slab
+				pushedWork = true
+			case twinHold != nil:
+				copy(twinHold[32:64], slab[0:32])
+				cp.hashSlab254(0, twinHold[0:64])
+				cp.layerQueues[myIdx+1] <- twinHold[0:32:64]
+				pushedWork = true
+				twinHold = nil
+			default:
+				twinHold = slab[0:32:64]
+			}
 
-				cp.hash254Into(cp.layerQueues[myIdx+1], chunkHold, chunk)
-				chunkHold = nil
+			// Check whether we need another worker
+			//
+			// n.b. we will not blow out of the preallocated layerQueues array,
+			// as we disallow Write()s above a certain threshold
+			if pushedWork && cp.layerQueues[myIdx+2] == nil {
+				cp.addLayer(myIdx + 1)
 			}
 		}
 	}()
 }
 
-func (cp *Calc) hash254Into(out chan<- []byte, half1ToOverwrite, half2 []byte) {
+func (cp *Calc) hashSlab254(layerIdx uint, slab []byte) {
 	h := shaPool.Get().(hash.Hash)
-	h.Reset()
-	h.Write(half1ToOverwrite)
-	h.Write(half2)
-	d := h.Sum(half1ToOverwrite[:0]) // callers expect we will reuse-reduce-recycle
-	d[31] &= 0x3F
-	out <- d
+
+	stride := 1 << (5 + layerIdx)
+	for i := 0; len(slab) > i+stride; i += 2 * stride {
+		h.Reset()
+		h.Write(slab[i : i+32])
+		h.Write(slab[i+stride : 32+i+stride])
+		h.Sum(slab[i:i])[31] &= 0x3F // callers expect we will reuse-reduce-recycle
+	}
+
 	shaPool.Put(h)
 }