Tree Optimization #120

claravanstaden · 2023-02-06T12:59:11Z

The current implementation of the FastSSZ tree fills empty nodes of the tree so that the number of tree leaves are to the power of 2. However, if the max capacity of a the tree leaves is very large (like BeaconState Validators, max capacity is 1 099 511 627 776), Go cannot preallocate such a large slice (see issue #119 (comment)). Even if it could allocate the slice, adding all the empty nodes isn't necessary and slows down the tree creation since it needs to populate all the empty leaves, plus their parents.

This PR changes the tree implementation to only add nodes that have values, and adds precomputed zero order hashes where the adjacent sibling requires a right-node sibling in order to hash the two nodes together.

Taking @protolambda's SSZ diagram as an example, in the current version of this library, nodes 7, 14 and 15 would be added as part of the tree. In this PR, only node 7 is added.

This improves the performance of GetTree significantly and enables the successful testing of all the spectests, where BeaconBlock, BeaconState and ExecutionPayload were previously skipped because of performance issues.

claravanstaden · 2023-02-06T13:00:21Z

go.mod

@@ -6,10 +6,14 @@ require (
 	github.com/golang/snappy v0.0.3
 	github.com/minio/sha256-simd v1.0.0
 	github.com/mitchellh/mapstructure v1.3.2
+	github.com/prysmaticlabs/gohashtree v0.0.1-alpha.0.20220714111606-acbb2962fb48
+	github.com/stretchr/testify v1.8.1


I added stretchr for testing assertions - it makes checking errors and equality a little simpler.

claravanstaden · 2023-02-06T13:05:16Z

sszgen/generator/hash.go

@@ -138,7 +138,7 @@ func (v *Value) hashTreeRoot(name string, appendBytes bool) string {
 			})
 		} else {
 			// dynamic bytes require special handling, need length mixed in
-			hMethod := "PutBytes"
+			hMethod := "Append"


The hash tree root of PreviousEpochParticipation and CurrentEpochParticipation does not work. Both PutBytes and MerkleizeWithMixin merkleizes the bytes if the bytes length is > 32, resulting in a double merkleization. The beacon spec tests doesn't uncover this bug because all the bytes in the spec tests for these two values are less than 32 bytes.

Can this happen in production? It might be because Prysm does not use fastssz for the BeaconState and this case is not covered.

Yes, it can happen in production. I believe that is correct. HashTreeRoot for the BeaconState crashed before it could get to this part of the code (it crashed at index 11, and this happens on index 15), so I think that's why we haven't encountered it before.

claravanstaden · 2023-02-06T13:07:28Z

hasher.go

+	if rest := len(input) % 32; rest != 0 {
+		// pad zero bytes to the left
+		input = append(input, zeroBytes[:32-rest]...)
+	}


Related to comment: https://github.com/ferranbt/fastssz/pull/120/files#r1097352609

PutBytes and MerkleizeWithMixin cannot be used together because they both merkleize bytes, resulting in double merkleization. The correct use is Append with MerkleizeWithMixin. However, MerkleizeWithMixin should then check if the bytes should be padded to the next 32 bytes increment.

Could you use FillUpTo32 instead at the beginning of the function for the whole h.buf array?

h.FillUpTo32() input := h.buf[indx:]

Good point! Addressed in b74a6c6.

claravanstaden · 2023-02-06T13:09:16Z

wrapper.go

+	// if byte list is empty, fill with zeros
+	if len(b) == 0 {
+		b = append(b, zeroBytes[:32]...)
+	}
+	// if byte list isn't filled with 32-bytes padded, pad
+	if rest := len(b) % 32; rest != 0 {
+		b = append(b, zeroBytes[:32-rest]...)
+	}


Similarly to https://github.com/ferranbt/fastssz/pull/120/files#r1097354675, we need to check if zero bytes need to added as padding to the next 32 byte increment.

claravanstaden · 2023-02-06T13:09:37Z

wrapper.go

-func (w *Wrapper) fillEmptyNodes(i int) {
+func (w *Wrapper) getLimit(i int) int {
 	size := len(w.nodes[i:])
-	for i := size; i < int(nextPowerOfTwo(uint64(size))); i++ {
-		w.nodes = append(w.nodes, EmptyLeaf())
-	}


We don't need to fill nodes anymore.

claravanstaden · 2023-02-06T13:11:03Z

spectests/beacon_state_bellatrix_test.go

+func TestBeaconStateTree_HashTreeRoot(t *testing.T) {
+	data, err := os.ReadFile(TestFileName)
+	require.NoError(t, err)
+
+	sszState := BeaconStateBellatrix{}
+	err = sszState.UnmarshalSSZ(data)
+	require.NoError(t, err)
+
+	tree, err := sszState.GetTree()
+	require.NoError(t, err)
+
+	hash := tree.Hash()
+
+	require.Equal(t, "c4a9c5ebf637c089db599574b568bb679b385c1984f08410707db08e03d7ae52", hex.EncodeToString(hash))
+}


These tests are a sanity check against real data, as some of the spec tests didn't catch some issues with the generated encodings.

claravanstaden · 2023-02-06T13:17:44Z

spectests/structs_test.go

-	if structName == "BeaconState" || structName == "BeaconBlockBody" || structName == "ExecutionPayload" {
-		// this gets to expensive, BeaconState even crashes with out-of-bounds memory allocation
-		return
-	}
-


This optimization allows running all the beacon spec tests.

claravanstaden · 2023-02-06T13:22:18Z

tree.go

+				// node with empty right node, add zero order hash as right node and mark right node as empty
+				if nodes[leftIndex] != nil && nodes[rightIndex] == nil {
+					nodes[i] = NewNodeWithLR(nodes[leftIndex], NewEmptyNode(zeroOrderHashes[k+1]))
+				}


Basically the crux of the optimization. 😄

claravanstaden · 2023-02-06T13:25:13Z

tree.go

+	leavesStart := powerTwo(depth)
+	leafIndex := numLeaves - 1
+
+	nodes := make(map[int]*Node)


Changes nodes to a map instead of a slice so that we can lookup child nodes using the computed left and right index.

claravanstaden · 2023-02-06T13:26:09Z

tree.go

-	right *Node
+	left    *Node
+	right   *Node
+	isEmpty bool


Adds a flag so we know the zero order hash is actually an empty node and we shouldn't traverse the tree further for the right node.

claravanstaden · 2023-02-06T13:42:06Z

spectests/structs_encoding.go

+		hh.Append(b.PreviousEpochParticipation)
 		hh.MerkleizeWithMixin(elemIndx, byteLen, (1099511627776+31)/32)


Prevents double merkleization of PutBytes and MerkleizeWithMixin used together.

ferranbt · 2023-02-06T16:36:24Z

This looks amazing! I will need a bit of time but It will most likely land by the end of the week. Could you regenerate the testcases (make generate-testcases) to pass the CI?

Check if removing go mod vendor works.

claravanstaden · 2023-02-06T18:42:23Z

Thanks @ferranbt, that sounds good. CI should pass now, I tested on our fork: https://github.com/Snowfork/fastssz/actions/runs/4106825072/jobs/7085540042

I had to add go mod vendor to the CI to download stretchr, not sure why it isn't necessary for the other go vendor dependencies.

claravanstaden · 2023-02-06T18:43:29Z

scripts/ci-validate-testcases.sh

@@ -5,7 +5,7 @@ make generate-testcases

 # check differences
 cd sszgen/testcases
-if [[ `git status --porcelain` ]]; then
+if [[ `git status --porcelain .` ]]; then


Only check the current (sszgen/testcases) directory, otherwise it picks up uncommitted changes in the rest of the project.

spectests/beacon_state_bellatrix_test.go

ferranbt

Approved! Some minor suggestions.

Co-authored-by: Ferran Borreguero <ferran.borreguero@gmail.com>

claravanstaden added 2 commits February 6, 2023 14:41

Tree optimizations.

2bdd995

Updates structs_encoding.go.

eeff6bf

claravanstaden commented Feb 6, 2023

View reviewed changes

claravanstaden marked this pull request as ready for review February 6, 2023 13:11

claravanstaden commented Feb 6, 2023

View reviewed changes

Adds a comment about test data.

9b5ac4b

claravanstaden commented Feb 6, 2023

View reviewed changes

ferranbt self-requested a review February 6, 2023 16:41

claravanstaden added 7 commits February 6, 2023 19:52

Update testcases.

602c5dd

touch

9043c09

Attempt to fix missing import

9de0c85

Check git status --porcelain output

1f21576

Updates Go version.

321a442

Check testcases directory only for testcases git status.

7de01f5

Check if removing go mod vendor works.

Add vendor dependencies.

72b5620

claravanstaden commented Feb 6, 2023

View reviewed changes

ferranbt reviewed Feb 7, 2023

View reviewed changes

spectests/beacon_state_bellatrix_test.go Outdated Show resolved Hide resolved

ferranbt approved these changes Feb 7, 2023

View reviewed changes

claravanstaden and others added 2 commits February 7, 2023 10:35

Update spectests/beacon_state_bellatrix_test.go

50b7a07

Co-authored-by: Ferran Borreguero <ferran.borreguero@gmail.com>

Minor improvements

b74a6c6

claravanstaden requested a review from ferranbt February 7, 2023 08:43

ferranbt merged commit dc27d62 into ferranbt:main Feb 7, 2023

ferranbt mentioned this pull request Feb 7, 2023

Running out of memory on BeaconStateBellatrix GetTree #119

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Tree Optimization #120

Tree Optimization #120

claravanstaden commented Feb 6, 2023 •

edited

Loading

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023 •

edited

Loading

ferranbt Feb 7, 2023

claravanstaden Feb 7, 2023

claravanstaden Feb 6, 2023 •

edited

Loading

ferranbt Feb 7, 2023 •

edited

Loading

claravanstaden Feb 7, 2023

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023

claravanstaden Feb 6, 2023

ferranbt commented Feb 6, 2023

claravanstaden commented Feb 6, 2023

claravanstaden Feb 6, 2023

ferranbt left a comment

		hh.Append(b.PreviousEpochParticipation)
		hh.MerkleizeWithMixin(elemIndx, byteLen, (1099511627776+31)/32)

Tree Optimization #120

Tree Optimization #120

Conversation

claravanstaden commented Feb 6, 2023 • edited Loading

Choose a reason for hiding this comment

claravanstaden Feb 6, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

claravanstaden Feb 6, 2023 • edited Loading

Choose a reason for hiding this comment

ferranbt Feb 7, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ferranbt commented Feb 6, 2023

claravanstaden commented Feb 6, 2023

Choose a reason for hiding this comment

ferranbt left a comment

Choose a reason for hiding this comment

claravanstaden commented Feb 6, 2023 •

edited

Loading

claravanstaden Feb 6, 2023 •

edited

Loading

claravanstaden Feb 6, 2023 •

edited

Loading

ferranbt Feb 7, 2023 •

edited

Loading