optimize CPU inference with Array-Based Tree Traversal

[razdoburdin]

This PR introduces optimization for CPU inference. For each tree, the top N levels are transformed into a compact array-based layout. This allows for a branchless node indexing rule: idx = 2 * idx + int(val < split_cond). To minimize memory overhead, this transformation from the standard tree structure to the array layout is performed on-the-fly for each block of data being processed. Even with this additional calculations, improved data locality in the cache-friendly array layout leads to inference speed up to ~2x (x1.4 on average).

[trivialfis]

Thank you for the optimization on the inference. Please unmark the "draft" status and ping me when the PR is ready for testing.

[Vika-F]

Cosmetic changes.

The next possible step would be to convert the trees into array-based representation only once, and not to do it for each block of data.

[razdoburdin]

The next possible step would be to convert the trees into array-based representation only once, and not to do it for each block of data.

it sounds reasonable and will further improve perf (by cost of increasing memory consumption).

[razdoburdin]

Thank you for the optimization on the inference. Please unmark the "draft" status and ping me when the PR is ready for testing.

hi @trivialfis, the PR is ready for review.

[trivialfis]

cc @hcho3

[trivialfis]

Still trying to understand the code, will give it a try later. In the meanwhile, could you please craft some specific unittests for the new inference algorithm?

[trivialfis]

What happens if this is a small online inference call? The input size could be a few samples per call.

[razdoburdin]

The default (the old one) realization will be used

[trivialfis]

unused?

[razdoburdin]

The optimized array_layout processing is effective only for nodes, that are close to root. For other nodes, we still use the original method.

[razdoburdin]

Still trying to understand the code, will give it a try later. In the meanwhile, could you please craft some specific unittests for the new inference algorithm?

I added some unit-tests.

[trivialfis]

I'm still trying to understand the code, in the meantime, let me do some refactoring in this and the next week to accommodate the new optimization. We need a better structure to handle all these:

• Predict with scalar leaf.
• Predict with vector leaf.
• Array predict with scalar leaf.
• Array predict with vector leaf.
• Column split with scalar leaf.

I think I will split up the CPU predictor into multiple pieces.

[trivialfis]

What happens if the tree is not well balanced and is more like a linked list than a tree?

[razdoburdin]

In this case we add dummy nodes with nan as split condition. In this dummy node the decision is always "go right" (see also comments https://github.com/razdoburdin/xgboost/blob/b0eaa856e1246416f7f9538bcc004e7723d9b997/src/predictor/array_tree_layout.h#L154), the left child are not initialized.

So with array layout we have to allocated all nodes, but keep some of the unpopulated in case the tree is pure balanced.

Initial tree:

Tree with dummy (nan-valued nodes):

[trivialfis]

Thank you for sharing, that makes sense.

[trivialfis]

Is there a benefit of doing this C++ overloading rather than the simpler tree.IsLeaf? How much faster are we seeing?

[razdoburdin]

I did the overload to handle both RegTree and MultiTargetTree cases. Is there a better option?

[trivialfis]

Use RegTree without extracting the Multi-target tree when populating the buffer, and delegate the dispatching to RegTree::LeftChild(bst_node_t nidx) instead of using the RegTree::Node::LeftChild. There's a check inside the RegTree::LeftChild:

  [[nodiscard]] bst_node_t LeftChild(bst_node_t nidx) const {
    if (IsMultiTarget()) {
      return this->p_mt_tree_->LeftChild(nidx);
    }
    return (*this)[nidx].LeftChild();
  }

[trivialfis]

I'm trying to cleanup the CPU predictor. I will update this PR once it is finished.

[trivialfis]

I need to fix a perf regression caused by the new ordinal encoder.

[trivialfis]

I need to fix a perf regression caused by the new ordinal encoder.

This has been fixed. I will look deeper into this PR.

[trivialfis]

Suggested change

	struct Empty>;
	Empty>;

[trivialfis]

Thank you for sharing, that makes sense.

[trivialfis]

Use RegTree without extracting the Multi-target tree when populating the buffer, and delegate the dispatching to RegTree::LeftChild(bst_node_t nidx) instead of using the RegTree::Node::LeftChild. There's a check inside the RegTree::LeftChild:

  [[nodiscard]] bst_node_t LeftChild(bst_node_t nidx) const {
    if (IsMultiTarget()) {
      return this->p_mt_tree_->LeftChild(nidx);
    }
    return (*this)[nidx].LeftChild();
  }

[trivialfis]

Thank you for expanding the tree layout. In the future (when you can prioritize it), do you think it's possible to create and store the layout inside the RegTree structure as an opt-in method call? My reasoning is as follows:

• The existing RegTree and the multi-target tree already use a very similar layout, minus the dummy nodes. It might be easier/cleaner to do it there.
• We can avoid complicating the predictor too much.
• We can cache the result in the regtree structure to avoid repeated initialization.

You can define a std::unique_ptr<ArrayTree> inside the RegTree, set it to nullptr. Define a method to create the array tree when needed, and reset it back to nullptr if any non-const method is called.