Advanced Systems Lab - Team #26

CPU-inference optimized variant of EzGATr (Easy Geometric Algebra Transformer).

Our optimizations target:

GenuineIntel, 11th Gen Core i7-1165G7, Tiger Lake (single thread)

Team members:

• Anastasija Tortevska
• Glenn Schönbächler
• Kristijan Zafirovski
• Mukhammadali Sayfiddinov

Slides: https://docs.google.com/presentation/d/1znu0kcxWRBG2q-T0RSiTfL3pvKltQKpec6rvvNRY7oQ/edit?usp=sharing

---

Repository overview

We optimized the hot CPU-inference kernels of EzGATr (geometric product, equi-join, equi-linear, RMS norm, gated GELU, and geometric attention) and rewrote them in C++ with hand-vectorized AVX2 / AVX-512 variants. The unmodified PyTorch library is the baseline; our work is additive and lives in a few places:

Where	What
src/ezgatr/_csrc/	All optimized kernels (C++/AVX). pga_ops.cpp (geometric product + join), attention_ops.cpp / attention_ops_v4.cpp, rms_ops.cpp, plus kernels/*.inc (per-strategy SIMD bodies) and bindings.cpp (pybind exports).
src/ezgatr/opt/	Python entry points exposing every kernel version to benchmarks and the model.
src/ezgatr/nn/functional/	Thin wiring so the existing model can call the optimized ops (e.g. attention_cpp.py).
benchmarks/	Benchmark runners, plotting scripts, and recorded results.
26_report.pdf	Final report.

Versioning. Each op carries numbered variants showing the optimization path: *_v0 / *_ver_0 is the straight scalar C++ port of the baseline, and higher numbers are successive steps (instruction-level parallelism, SoA + AVX2, register-resident, AVX-512). The unsuffixed name (e.g. geometric_product) is the selected best variant used in the end-to-end model.

Diff against the upstream port. Everything after the original code was ported is ours:

git diff 81cd4f1 -- src/ezgatr   # 81cd4f1 = "all code from original ezgatr is ported"

Build & reproduce. Build the extension with python setup.py build_ext --inplace, then follow benchmarks/README.md for the benchmark and roofline commands. All numbers are single-thread on a Tiger Lake i7-1165G7.

---

Below is the original README for the project.

Geometric Algebra Transformer Made Easy

What is EzGATr?

EzGATr (Easy Geometric Algebra Transformer) intends to be a simple-to-use and lightweight Python library for building 3D Geometric Algebra Transformers (GATr). It is a collection of operators, modules, utilities, etc. build on top of PyTorch. In addition, EzGATr also seeks to bridge the gap between the mathematical formulations and corresponding implementations through extensive documentation and explanations to facilitate learning and potential future optimizations.

Installation

EzGATr is currently under active development. To install it, you can clone the repository and install it using pip. Use the -e flag to install it in editable mode for making quick changes.

$ git clone https://github.com/Guest400123064/ezgatr.git
$ cd ezgatr
$ pip install -e .

Alternatively, one can directly install it from PyPI.

$ pip install ezgatr

Building the C++ extension

The optimized operators live in ezgatr.opt._opt_ops and need to be rebuilt after changing the sources under src/ezgatr/_csrc/.

$ python setup.py build_ext --inplace

Usage

The quick start example shown below demonstrates how to use EzGATr to build a equivariant network with one input/output channel (e.g., a point cloud) and four hidden channels. The mock input data contains a batch of eight samples, each with 256 3D objects embedded as multi-vectors.

import torch
import torch.nn as nn

from ezgatr.nn import EquiLinear
from ezgatr.nn.functional import scaler_gated_gelu


class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = EquiLinear(1, 4)
        self.fc2 = EquiLinear(4, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = scaler_gated_gelu(x)
        x = self.fc2(x)
        return x


net = SimpleNet()
in_ = torch.randn(8, 256, 1, 16)
out = net(in_)

One can refer to this example for how to build a full-fledged GATr model with EzGATr, involving equivariant geometric attention and geometric MLP.

API References

The complete API references for EzGATr can be found here. TL;DR, the package is organized as follows:

• ezgatr.nn: Contains the core modules and layers for building GATr models. It is organized similarly to PyTorch's torch.nn package, where the functional submodule contains lower-level operators and transformations.
• ezgatr.interfaces: Contains the utility functions that help encode and decode 3D objects to and from multi-vectors.
• ezgatr.nets: Contains off-of-the-shelf networks built with EzGATr building blocks. It can also be used as references for building custom networks with EzGATr.

Citation and Authors

If you find EzGATr useful in your research, please consider citing it using the following BibTeX entry:

@misc{ezgatr2024,
  author = {Yuxuan Wang, Xiatao Sun},
  title = {EzGATr: Geometric Algebra Transformer Made Easy},
  year = 2024,
  publisher = {Zenodo},
  version = {v0.1.0-alpha},
  doi = {10.5281/zenodo.13920438},
  url = {https://doi.org/10.5281/zenodo.13920438}
}

EzGATr leveraged the work from GATr (official GitHub repository) by Johann Brehmer, Pim de Haan, Sönke Behrends, and Taco Cohen, published at NeurIPS 2023. Please consider citing their work as well.

@inproceedings{NEURIPS2023_6f6dd92b,
  author = {Brehmer, Johann and de Haan, Pim and Behrends, S\"{o}nke and Cohen, Taco S},
  booktitle = {Advances in Neural Information Processing Systems},
  editor = {A. Oh and T. Naumann and A. Globerson and K. Saenko and M. Hardt and S. Levine},
  pages = {35472--35496},
  publisher = {Curran Associates, Inc.},
  title = {Geometric Algebra Transformer},
  url = {https://proceedings.neurips.cc/paper_files/paper/2023/file/6f6dd92b03ff9be7468a6104611c9187-Paper-Conference.pdf},
  volume = {36},
  year = {2023}
}

License

EzGATr is distributed under the terms of the MIT license.