PyTorch-based Geometric Learning with Non-CUDA Processing Units: Experiences from Intel Gaudi-v2 HPUs

TL;DR

This paper details the adaptation of PyTorch-based geometric learning frameworks to Intel Gaudi-v2 HPUs, providing utilities, tutorials, and examples to facilitate research on non-CUDA hardware.

Contribution

It introduces utilities and comprehensive resources for porting geometric learning to Gaudi-v2 HPUs, enhancing cross-platform experimentation and development.

Findings

Restored essential geometric operations on Gaudi-v2 HPUs

Compiled tutorials and real-world examples with diagnostic analyses

Published a GitHub repository to support community efforts

Abstract

Geometric learning has emerged as a powerful paradigm for modeling non-Euclidean data, especially graph-structured ones, with applications spanning social networks, molecular structures, knowledge graphs, and recommender systems. While Nvidia's CUDA-enabled graphics processing units (GPUs) largely dominate the hardware landscape, emerging accelerators such as Intel's Gaudi Habana Processing Units (HPUs) offer competitive performance and energy efficiency. However, the usage of such non-CUDA processing units requires significant engineering effort and novel software adaptations. In this work, we present our experiences porting PyTorch-based geometric learning frameworks to Gaudi-v2 HPUs. We introduce a collection of core utilities that restore essential operations (e.g., scatter, sparse indexing, k-nearest neighbors) on Gaudi-v2 HPUs, and we consolidate sixteen guided tutorials and eleven real-world examples with diagnostic analyses of encountered failures and detailed workarounds. We collect all our experiences into a publicly accessible GitHub repository. Our contributions lower the barrier for researchers to experiment with geometric-learning algorithms and models on non-CUDA hardware, providing a foundation for further optimization and cross-platform portability.