| dc.description.abstract | The proliferation of machine learning together with the rapid evolution of the hardware ecosystem has led to a surge in the demand for model inference on a variety of hardware. Decision tree based models are the most popular models on tabular data. Our work is motivated by the problems encountered when targeting inference of these models to run at peak performance on CPU and GPU targets. Decision tree ensemble inference is usually performed with libraries such as XGBoost, LightGBM, and Sklearn on CPUs and RAPIDs FIL on the GPU. These libraries incorporate a fixed set of optimizations for the hardware targets they support. However, these solutions are neither portable nor achieve the best possible performance for the specific hardware they target. This is because they do not specialize the inference code to the model being used and the specific target processor, leaving significant performance on the table.
To address this problem, we designed and built TREEBEARD, a schedule-guided, retargetable compiler for decision tree based models that searches over several optimization choices and automatically generates high-performance inference routines for CPUs and GPUs. TREEBEARD has two core components. The first is a scheduling language that encapsulates the optimization space, and techniques to efficiently explore this space. The second is an optimizing retargetable compiler that can generate code for any specified schedule by lowering the inference computation to optimized CPU or GPU code through multiple intermediate abstractions. By applying model-specific optimizations at the higher levels, tree walk optimizations at the middle level, and machine-specific optimizations lower down, TREEBEARD can specialize inference code for each model on each supported target. TREEBEARD combines several novel optimizations at various abstraction levels, uses different data layouts, loop structures and caching strategies to mitigate architectural bottlenecks and achieve portable performance across a range of targets.
We implement TREEBEARD using the MLIR compiler infrastructure and demonstrate its utility by evaluating it on a diverse set of benchmarks. TREEBEARD is significantly faster on the CPU than state-of-the-art systems, XGBoost, Treelite and HUMMINGBIRD, by 2.6x, 4.7x and 5.4x respectively in a single-core execution setting, and by 2.3x, 2.7x and 14x respectively in multi-core settings. On the GPU, TREEBEARD generated code is an order of magnitude faster than XGBoost and about 2-5x faster on average than RAPIDS FIL and Tahoe over several batch sizes. While these systems only target NVIDIA GPUs, TREEBEARD achieves competent performance on AMD GPUs as well.
Overall, this thesis demonstrates that a domain-specific, schedule-guided compiler can provide portable, high-performance inference for decision tree ensembles, significantly outperforming hand-optimized libraries and other more limited compilers across a diverse set of models and hardware targets. | en_US |
