Paper Benchmark Routines ?

[mitkotak]

I know you guys are planning to do this but just wanted to make note of it. Would appreciate if there's some set of scripts that I can trigger overnight to get the paper plots on my RTX A5500. Thanks !

[asglover]

This is how to run the benchmarks from which the plots are generated.

https://github.com/vbharadwaj-bk/OpenEquivariance/blob/main/test/benchmark_driver.py

We can follow up with more details about how to generate the plots, but first let’s see that the benchmarks will run.

[vbharadwaj-bk]

We have started providing a CLI interface to run benchmarks, which will be rolled out in a sequence of PRs (issue will be closed when all tests are runnable). test/benchmark.py is now runnable from the command line with options to run the "UVU" and roofline benchmarks - more coming soon.

These two commands will run the benchmarks as we used for our paper:

python test/benchmark.py -o <OUTPUT_FOLDER> uvu
python test/benchmark.py -o <OUTPUT_FOLDER> roofline

For finer control and information about more command line arguments, append --help after the benchmark name.

For the A550 and the uvu benchmark, you will likely need to decrease the batch size from 50k to something manageable, or try running the benchmarks without the e3nn implementation which consumes more memory. Both options available from CLI.

[vbharadwaj-bk]

@mitkotak Hopefully #65 resolves this issue; our benchmark scripts are tailored for the high main memory / shared memory of the A100, but we do have some options to get plots on a GPU with lower memory. Here's a set of invocations that I tested on an A5000 (see the README.md for more details on the packages you need to run this):

python benchmark.py -o outputs/uvu uvu --limited-memory --plot
python benchmark.py -o outputs/uvw uvw -b 25000 --plot
python benchmark.py -o outputs/roofline roofline --plot
python benchmark.py -o outputs/conv conv --data data/molecular_structures --limited-memory

Let us know if you can run these / what issues you face, and we can get this closed.

[mitkotak]

Have you seen this build error ?

 make -j4
[ 50%] Building CXX object CMakeFiles/kernel_wrapper.dir/kernel_wrapper.cpp.o
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:17:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:54:5: error: ‘KernelLaunchConfig’ does not name a type
   54 |     KernelLaunchConfig forward_config;
      |     ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:55:5: error: ‘KernelLaunchConfig’ does not name a type
   55 |     KernelLaunchConfig backward_config;
      |     ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:59:9: error: ‘KernelLaunchConfig’ has not been declared
   59 |         KernelLaunchConfig &forward_config_i,
      |         ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:60:9: error: ‘KernelLaunchConfig’ has not been declared
   60 |         KernelLaunchConfig &backward_config_i) :
      |         ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp: In constructor ‘JITTPImpl<JIT_IMPL>::JITTPImpl(std::string, int&, int&)’:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:62:13: error: class ‘JITTPImpl<JIT_IMPL>’ does not have any field named ‘forward_config’
   62 |             forward_config(forward_config_i),
      |             ^~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:63:13: error: class ‘JITTPImpl<JIT_IMPL>’ does not have any field named ‘backward_config’
   63 |             backward_config(backward_config_i) {
      |             ^~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:67:12: error: ‘forward_config’ was not declared in this scope; did you mean ‘forward_config_i’?
   67 |         if(forward_config.smem > 0) {
      |            ^~~~~~~~~~~~~~
      |            forward_config_i
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:71:12: error: ‘backward_config’ was not declared in this scope; did you mean ‘backward_config_i’?
   71 |         if(backward_config.smem > 0) {
      |            ^~~~~~~~~~~~~~~
      |            backward_config_i
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp: In member function ‘void JITTPImpl<JIT_IMPL>::exec_tensor_product(uint64_t, void*, void*, void*, void*)’:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:84:30: error: ‘forward_config’ was not declared in this scope
   84 |         jit.execute(0, args, forward_config);
      |                              ^~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp: In member function ‘void JITTPImpl<JIT_IMPL>::backward(size_t, void*, void*, void*, void*, void*, void*, void*)’:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/tensorproducts.hpp:94:30: error: ‘backward_config’ was not declared in this scope
   94 |         jit.execute(1, args, backward_config);
      |                              ^~~~~~~~~~~~~~~
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:18:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp: At global scope:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:98:5: error: ‘KernelLaunchConfig’ does not name a type
   98 |     KernelLaunchConfig forward_config;
      |     ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:99:5: error: ‘KernelLaunchConfig’ does not name a type
   99 |     KernelLaunchConfig backward_config;
      |     ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:103:9: error: ‘KernelLaunchConfig’ has not been declared
  103 |         KernelLaunchConfig forward_config_i,
      |         ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:104:9: error: ‘KernelLaunchConfig’ has not been declared
  104 |         KernelLaunchConfig backward_config_i) :
      |         ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp: In constructor ‘JITConvImpl<JIT_IMPL>::JITConvImpl(std::string, int, int)’:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:106:13: error: class ‘JITConvImpl<JIT_IMPL>’ does not have any field named ‘forward_config’
  106 |             forward_config(forward_config_i),
      |             ^~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:107:13: error: class ‘JITConvImpl<JIT_IMPL>’ does not have any field named ‘backward_config’
  107 |             backward_config(backward_config_i) {
      |             ^~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:112:12: error: ‘forward_config’ was not declared in this scope; did you mean ‘forward_config_i’?
  112 |         if(forward_config.smem > 0) {
      |            ^~~~~~~~~~~~~~
      |            forward_config_i
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:116:12: error: ‘backward_config’ was not declared in this scope; did you mean ‘backward_config_i’?
  116 |         if(backward_config.smem > 0) {
      |            ^~~~~~~~~~~~~~~
      |            backward_config_i
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp: In member function ‘void JITConvImpl<JIT_IMPL>::exec_conv(void*, void*, void*, void*, void*, void*, uint64_t, uint64_t, void*)’:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:135:30: error: ‘forward_config’ was not declared in this scope
  135 |         jit.execute(0, args, forward_config);
      |                              ^~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:140:13: error: ‘KernelLaunchConfig’ was not declared in this scope
  140 |             KernelLaunchConfig fixup_config;
      |             ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:141:13: error: ‘fixup_config’ was not declared in this scope
  141 |             fixup_config.num_blocks = forward_config.num_blocks;
      |             ^~~~~~~~~~~~
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:18:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp: In member function ‘void JITConvImpl<JIT_IMPL>::backward(void*, void*, void*, void*, void*, void*, void*, void*, void*, uint64_t, uint64_t, void*, void*)’:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:161:30: error: ‘backward_config’ was not declared in this scope
  161 |         jit.execute(1, args, backward_config);
      |                              ^~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:166:13: error: ‘KernelLaunchConfig’ was not declared in this scope
  166 |             KernelLaunchConfig fixup_config;
      |             ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/convolution.hpp:167:13: error: ‘fixup_config’ was not declared in this scope
  167 |             fixup_config.num_blocks = backward_config.num_blocks;
      |             ^~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp: In function ‘void pybind11_init_kernel_wrapper(pybind11::module_&)’:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:28:26: error: ‘JITKernel’ was not declared in this scope
   28 |     py::class_<JITTPImpl<JITKernel>, GenericTensorProductImpl>(m, "JITTPImpl")
      |                          ^~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:28:35: error: template argument 1 is invalid
   28 |     py::class_<JITTPImpl<JITKernel>, GenericTensorProductImpl>(m, "JITTPImpl")
      |                                   ^
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:28:62: error: template argument 1 is invalid
   28 |     py::class_<JITTPImpl<JITKernel>, GenericTensorProductImpl>(m, "JITTPImpl")
      |                                                              ^
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:36: error: ‘KernelLaunchConfig’ was not declared in this scope
   29 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |                                    ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:18: error: parse error in template argument list
   29 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |                  ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:77: error: no matching function for call to ‘init<std::string, <expression error>, <expression error> >()’
   29 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:1:
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1958:40: note: candidate: ‘template<class ... Args> pybind11::detail::initimpl::constructor<Args ...> pybind11::init()’
 1958 | detail::initimpl::constructor<Args...> init() {
      |                                        ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1958:40: note:   template argument deduction/substitution failed:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:77: error: template argument 2 is invalid
   29 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:77: error: template argument 3 is invalid
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:1:
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1970:5: note: candidate: ‘template<class Func, class Ret> Ret pybind11::init(Func&&)’
 1970 | Ret init(Func &&f) {
      |     ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1970:5: note:   template argument deduction/substitution failed:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:77: error: wrong number of template arguments (3, should be at least 1)
   29 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:1:
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1970:5: note: provided for ‘template<class Func, class Ret> Ret pybind11::init(Func&&)’
 1970 | Ret init(Func &&f) {
      |     ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1978:5: note: candidate: ‘template<class CFunc, class AFunc, class Ret> Ret pybind11::init(CFunc&&, AFunc&&)’
 1978 | Ret init(CFunc &&c, AFunc &&a) {
      |     ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1978:5: note:   template argument deduction/substitution failed:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:77: error: template argument 2 is invalid
   29 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:29:77: error: template argument 3 is invalid
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:35:55: error: template argument 1 is invalid
   35 |     py::class_<JITConvImpl<JITKernel>, ConvolutionImpl>(m, "JITConvImpl")
      |                                                       ^
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:36:18: error: parse error in template argument list
   36 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |                  ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:36:77: error: no matching function for call to ‘init<std::string, <expression error>, <expression error> >()’
   36 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:1:
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1958:40: note: candidate: ‘template<class ... Args> pybind11::detail::initimpl::constructor<Args ...> pybind11::init()’
 1958 | detail::initimpl::constructor<Args...> init() {
      |                                        ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1958:40: note:   template argument deduction/substitution failed:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:36:77: error: template argument 2 is invalid
   36 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:36:77: error: template argument 3 is invalid
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:1:
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1970:5: note: candidate: ‘template<class Func, class Ret> Ret pybind11::init(Func&&)’
 1970 | Ret init(Func &&f) {
      |     ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1970:5: note:   template argument deduction/substitution failed:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:36:77: error: wrong number of template arguments (3, should be at least 1)
   36 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
In file included from /home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:1:
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1970:5: note: provided for ‘template<class Func, class Ret> Ret pybind11::init(Func&&)’
 1970 | Ret init(Func &&f) {
      |     ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1978:5: note: candidate: ‘template<class CFunc, class AFunc, class Ret> Ret pybind11::init(CFunc&&, AFunc&&)’
 1978 | Ret init(CFunc &&c, AFunc &&a) {
      |     ^~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/.venv/include/pybind11/pybind11.h:1978:5: note:   template argument deduction/substitution failed:
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:36:77: error: template argument 2 is invalid
   36 |         .def(py::init<std::string, KernelLaunchConfig&, KernelLaunchConfig&>());
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:36:77: error: template argument 3 is invalid
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:41:39: error: ‘KernelLaunchConfig’ is not a class, namespace, or enumeration
   41 |         .def_readwrite("num_blocks", &KernelLaunchConfig::num_blocks)
      |                                       ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:42:38: error: ‘KernelLaunchConfig’ is not a class, namespace, or enumeration
   42 |         .def_readwrite("warp_size", &KernelLaunchConfig::warp_size)
      |                                      ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:43:40: error: ‘KernelLaunchConfig’ is not a class, namespace, or enumeration
   43 |         .def_readwrite("num_threads", &KernelLaunchConfig::num_threads)
      |                                        ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:44:33: error: ‘KernelLaunchConfig’ is not a class, namespace, or enumeration
   44 |         .def_readwrite("smem", &KernelLaunchConfig::smem);
      |                                 ^~~~~~~~~~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:46:16: error: ‘DeviceProp’ was not declared in this scope
   46 |     py::class_<DeviceProp>(m, "DeviceProp")
      |                ^~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:46:26: error: template argument 1 is invalid
   46 |     py::class_<DeviceProp>(m, "DeviceProp")
      |                          ^
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:48:32: error: ‘DeviceProp’ is not a class, namespace, or enumeration
   48 |         .def_readonly("name", &DeviceProp::name)
      |                                ^~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:49:36: error: ‘DeviceProp’ is not a class, namespace, or enumeration
   49 |         .def_readonly("warpsize", &DeviceProp::warpsize)
      |                                    ^~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:50:33: error: ‘DeviceProp’ is not a class, namespace, or enumeration
   50 |         .def_readonly("major", &DeviceProp::major)
      |                                 ^~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:51:33: error: ‘DeviceProp’ is not a class, namespace, or enumeration
   51 |         .def_readonly("minor", &DeviceProp::minor)
      |                                 ^~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:52:47: error: ‘DeviceProp’ is not a class, namespace, or enumeration
   52 |         .def_readonly("multiprocessorCount", &DeviceProp::multiprocessorCount)
      |                                               ^~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:53:48: error: ‘DeviceProp’ is not a class, namespace, or enumeration
   53 |         .def_readonly("maxSharedMemPerBlock", &DeviceProp::maxSharedMemPerBlock);
      |                                                ^~~~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:55:31: error: ‘Allocator’ was not declared in this scope; did you mean ‘alloca’?
   55 |     py::class_<PyDeviceBuffer<Allocator>>(m, "DeviceBuffer")
      |                               ^~~~~~~~~
      |                               alloca
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:55:31: error: template argument 1 is invalid
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:55:40: error: template argument 1 is invalid
   55 |     py::class_<PyDeviceBuffer<Allocator>>(m, "DeviceBuffer")
      |                                        ^~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:61:16: error: ‘GPUTimer’ was not declared in this scope
   61 |     py::class_<GPUTimer>(m, "GPUTimer")
      |                ^~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:61:24: error: template argument 1 is invalid
   61 |     py::class_<GPUTimer>(m, "GPUTimer")
      |                        ^
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:63:24: error: ‘GPUTimer’ is not a class, namespace, or enumeration
   63 |         .def("start", &GPUTimer::start)
      |                        ^~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:64:41: error: ‘GPUTimer’ is not a class, namespace, or enumeration
   64 |         .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed)
      |                                         ^~~~~~~~
/home/mkotak/atomic_architects/projects/OpenEquivariance/openequivariance/extension/kernel_wrapper.cpp:65:33: error: ‘GPUTimer’ is not a class, namespace, or enumeration
   65 |         .def("clear_L2_cache", &GPUTimer::clear_L2_cache);
      |                                 ^~~~~~~~
make[2]: *** [CMakeFiles/kernel_wrapper.dir/build.make:76: CMakeFiles/kernel_wrapper.dir/kernel_wrapper.cpp.o] Error 1
make[1]: *** [CMakeFiles/Makefile2:83: CMakeFiles/kernel_wrapper.dir/all] Error 2
make: *** [Makefile:136: all] Error 2

[vbharadwaj-bk]

Ah is this an earlier version of the package? We switched away from CMakeLists entirely; you shouldn't need to manually build the extension anymore, a fresh pip install should work. But that is a good point we should probably update or deprecate the devbuild instructions.

[mitkotak]

In the benchmarking machinery, is there a convenient place to put a torch.testing.assert between all the outputs ?

[asglover]

We can probably add a few lines in the existing correctness checking machinery. Let me take a look.

Have you been able to run the benchmarks successfully?
What is the purpose of using the torch.testing.assert? For some correctness checking automation?

Remember that there is a weight re-ordering step required for some tensor products.

[mitkotak]

We can probably add a few lines in the existing correctness checking machinery. Let me take a look.

Have you been able to run the benchmarks successfully?

Still running through it.

What is the purpose of using the torch.testing.assert? For some correctness checking automation?

Yup

Remember that there is a weight re-ordering step required for some tensor products.

Ohh so is there a transpose at the end to get the outputs right ?

[asglover]

The weights must be reordered so that that function produces the same outputs.

OpenEquivariance/openequivariance/implementations/ComputationSchedule.py

Line 427 in 8899985

def reorder_weights(self, weights_in, weights_out, direction, has_batch_dim):

If you wanted to get the same results.
You'd initialize the TensorProduct
Which would initialize a ComputationSchedule inside the TensorProduct
You can use the re-order weights function associated with the ComputationSchedule to reorder the weights.
Then you call the tensor product and should get the same outputs.

In many cases, the weight reordering is a no-op.
How you want to use it depends on the use case. For checking correctness, you might want to run it even time because performance is not as critical, but it production, you would want to statically reorder the weights once. For systems like nequip, you can reorder the weights in the dense NN which creates the weights which go into the Tensor Product. This is the same as reordering the weights with the function in practice, but as a one time cost.

[vbharadwaj-bk]

To add to this @mitkotak: no transpose is needed for the output or input (only the weights, sometimes), and I can confirm for Nequip, MACE, and any channel wise tensorproduct with SH multiplicity 1, it is a no-op. You only need to reorder weights for certain uvw configs, or if the spherical harmonics have multiplicity greater than 1.

[vbharadwaj-bk]

Ah yes, to answer your question directly: you can insert an assertion just after line 138 of TensorProductBase.py to test correctness the forward pass, and likewise for the backward pass: https://github.com/PASSIONLab/OpenEquivariance/blob/main/openequivariance/implementations/TensorProductBase.py#L138. If you'd rather rely on our correctness machinery, you can set correctness=True in benchmark.py and we will check against e3nn for you.

[asglover]

@mitkotak
#61 (comment)

That is interesting. It could be implementation details. Were you using their preferred native memory layout, or did you include a transpose from the e3nn native memory layout? If there is a difference, it could be from the layout or other details.

I will take a look at this and get back to you. Do you still have NERSC access? Have you tried running these benchmarks on the paper hardware?

[mitkotak]

@mitkotak #61 (comment)

That is interesting. It could be implementation details. Were you using their preferred native memory layout, or did you include a transpose from the e3nn native memory layout? If there is a difference, it could be from the layout or other details.

This is without the overhead of TransposeIrrepsLayout

I will take a look at this and get back to you. Do you still have NERSC access? Have you tried running these benchmarks on the paper hardware?

Nope.