MACE Integration Utilities (#96)

* Added torch datatype converter function for convenience.

* Added weight_numel to torch wrappers.

* Minor fix to TensorProduct.

* Reverted mace_driver for now.

* Updated README.

* Updated README with instructions.

* Updated the README one more time.

Author: vbharadwaj-bk

.gitignore

@@ -25,6 +25,7 @@ nvidia-mathdx*
 .vscode/*
 *.ncu-rep
 mace_dev
+mace_oeq_integration
 valid_indices*
 
 *.xyz

README.md

@@ -4,23 +4,25 @@
 [[Supported Tensor Products]](#tensor-products-we-accelerate)
 [[Citation and Acknowledgements]](#citation-and-acknowledgements)
 
-OpenEquivariance is a kernel generator for the Clebsch-Gordon tensor product, 
+OpenEquivariance is a CUDA and HIP kernel generator for the Clebsch-Gordon tensor product, 
 a key kernel in rotation-equivariant deep neural networks. 
 It implements some of the tensor products 
-that [e3nn](https://e3nn.org/) supports that are
+that [e3nn](https://e3nn.org/) supports 
 commonly found in graph neural networks 
 (e.g. [Nequip](https://github.com/mir-group/nequip) or
-[MACE](https://github.com/ACEsuit/mace)). To get started, install our package via
+[MACE](https://github.com/ACEsuit/mace)). To get 
+started, ensure that you have GCC 9+ on your system 
+and install our package via
 
 ```bash
 pip install git+https://github.com/PASSIONLab/OpenEquivariance
 ```
 
-We provide up to an order of magnitude acceleration over e3nn
-and up to ~2x speedup over 
-[NVIDIA cuEquivariance](https://github.com/NVIDIA/cuEquivariance),
-which has a closed-source kernel package. We also offer fused
-equivariant graph convolutions that can reduce 
+We provide up to an order of magnitude acceleration over e3nn perform on par with the latest
+version of [NVIDIA cuEquivariance](https://github.com/NVIDIA/cuEquivariance),
+which has a closed-source kernel package. 
+We also offer fused equivariant graph 
+convolutions that can reduce 
 computation and memory consumption significantly. 
 
 We currently support NVIDIA GPUs and just added beta support on AMD GPUs for
@@ -124,22 +126,24 @@ print(torch.norm(Z))
 arbitrary order. 
 
 ## Installation 
-We currently support Linux systems only. We recommend that you use
-`conda` or `mamba` to set up a Python environment for installation.
-
-After activating an environment of your choice, run
+We currently support Linux systems only. 
+Before installation and the first library import, 
+ensure that the command 
+`c++ --version` returns GCC 9+; if not, set the
+`CC` and `CXX` environment variables to point to
+valid compilers. On NERSC Perlmutter,
+`module load gcc` will set up your environment
+correctly. 
+
+To install, run
 ```bash
 pip install git+https://github.com/PASSIONLab/OpenEquivariance
 ```
 After installation, the very first library
-import will trigger a build of a C++ extension we use.
+import will trigger a build of a C++ extension we use,
+which takes longer than usual.
 All subsequent imports will not retrigger compilation.
 
-If you encounter problems with installation, let us
-know by filing a bug and try a development build (see
-below). After installation, you should be able 
-to run the example above.
-
 ## Replicating our benchmarks 
 To run our benchmark suite, you'll also need the following packages: 
 - `e3nn`, 

openequivariance/__init__.py

@@ -5,6 +5,7 @@
 from openequivariance.implementations.e3nn_lite import TPProblem, Irreps
 from openequivariance.implementations.TensorProduct import TensorProduct 
 from openequivariance.implementations.convolution.TensorProductConv import TensorProductConv 
+from openequivariance.implementations.utils import torch_to_oeq_dtype
 
 __version__ = version("openequivariance")
 

openequivariance/implementations/TensorProduct.py

@@ -10,6 +10,7 @@ class TensorProduct(torch.nn.Module, LoopUnrollTP):
     def __init__(self, problem, torch_op=True):
         torch.nn.Module.__init__(self)
         LoopUnrollTP.__init__(self, problem, torch_op)
+        self.weight_numel = problem.weight_numel
 
     @staticmethod
     def name():

openequivariance/implementations/convolution/TensorProductConv.py

@@ -15,6 +15,7 @@ def __init__(self, config, idx_dtype=np.int64, torch_op=True, deterministic=Fals
                 torch_op=torch_op, deterministic=deterministic)
 
         self.dummy_transpose_perm = torch.zeros(1, dtype=torch.int64, device='cuda')
+        self.weight_numel = self.config.weight_numel
 
         if extlib.TORCH_COMPILE:
             self.forward = self.forward_deterministic if deterministic else self.forward_atomic

openequivariance/implementations/utils.py

@@ -70,4 +70,18 @@ def filter_and_analyze_problem(problem):
     result = {
         "is_uvw": problem.instructions[0].connection_mode == "uvw", 
     }
-    return result
+    return result
+
+def torch_to_oeq_dtype(torch_dtype):
+    global torch
+    import torch
+
+    """
+    Converts torch dtype to oeq dtype
+    """
+    if torch_dtype == torch.float32:
+        return np.float32
+    elif torch_dtype == torch.float64:
+        return np.float64
+    else:
+        raise ValueError("Unsupported torch dtype!")