{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/an-efficient-sparse-kernel-generator-for-o-3","title":"An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks","arxiv_id":"2501.13986","date":"2025-01-23","proceeding":null,"authors":["Vivek Bharadwaj","Austin Glover","Aydin Buluc","James Demmel"],"abstract":"Rotation equivariant graph neural networks, i.e. networks designed to guarantee certain geometric relations between their inputs and outputs, yield state of the art performance on spatial deep learning tasks. They exhibit high data efficiency during training and significantly reduced inference time for interatomic potential calculations compared to classical approaches. Key to these models is the Clebsch-Gordon (CG) tensor product, a kernel that contracts two dense feature vectors with a highly-structured sparse tensor to produce a dense output vector. The operation, which may be repeated millions of times for typical equivariant models, is a costly and inefficient bottleneck. We introduce a GPU sparse kernel generator for the CG tensor product that provides significant speedups over the best existing open and closed-source implementations. Our implementation achieves high performance by carefully managing the limited GPU shared memory through static analysis at model compile-time, minimizing reads and writes to global memory. We break the tensor product into a series of smaller kernels with operands that fit entirely into registers, enabling us to emit long arithmetic instruction streams that maximize instruction-level parallelism. By fusing the CG tensor product with a subsequent graph convolution, we reduce both intermediate storage and global memory traffic over naive approaches that duplicate input data. We also provide optimized kernels for the gradient of the CG tensor product and a novel identity for the higher partial derivatives required to predict interatomic forces. Our kernels offer up to 1.3x speedup over NVIDIA's closed-source cuEquivariance package, as well as 10x speedup over the widely-used e3nn package. In FP64 precision, we offer up to 6.2x inference-time speedup for the MACE chemistry foundation model over the original unoptimized version.","url_abs":"https://arxiv.org/abs/2501.13986v4","url_pdf":"https://arxiv.org/pdf/2501.13986v4.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"an-efficient-sparse-kernel-generator-for-o-3","repo_url":"https://github.com/passionlab/openequivariance","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"BSD-3-Clause"}},{"paper_slug":"an-efficient-sparse-kernel-generator-for-o-3","repo_url":"https://github.com/vbharadwaj-bk/OpenEquivariance","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"BSD-3-Clause"}}],"tasks":[{"task_slug":null,"task_name":"GPU"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2501.13986","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2501.13986"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/passionlab/openequivariance","reach":{"status":"ok","spdx":"BSD-3-Clause"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/vbharadwaj-bk/OpenEquivariance","reach":{"status":"ok","spdx":"BSD-3-Clause"}}],"summary":{"unverified":7},"by_repo_kind":{"official":{"samples":7,"ran":0,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"b4495e5b43c77482","entry":"change_basis_real_to_complex","repo":"passionlab/openequivariance","repo_kind":"official","path":"openequivariance/openequivariance/core/e3nn_lite.py","file_url":"https://github.com/passionlab/openequivariance/blob/HEAD/openequivariance/openequivariance/core/e3nn_lite.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"b4495e5b43c77482"}},{"code_sha256_prefix":"df4032bb75e2f50e","entry":"create_schedule_case2","repo":"passionlab/openequivariance","repo_kind":"official","path":"openequivariance/openequivariance/core/ComputationSchedule.py","file_url":"https://github.com/passionlab/openequivariance/blob/HEAD/openequivariance/openequivariance/core/ComputationSchedule.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"df4032bb75e2f50e"}},{"code_sha256_prefix":"c4b2e015b2d92b98","entry":"create_schedule_case3","repo":"passionlab/openequivariance","repo_kind":"official","path":"openequivariance/openequivariance/core/ComputationSchedule.py","file_url":"https://github.com/passionlab/openequivariance/blob/HEAD/openequivariance/openequivariance/core/ComputationSchedule.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"c4b2e015b2d92b98"}},{"code_sha256_prefix":"074719fb4fa857ab","entry":"perm_inverse","repo":"passionlab/openequivariance","repo_kind":"official","path":"openequivariance/openequivariance/core/e3nn_lite.py","file_url":"https://github.com/passionlab/openequivariance/blob/HEAD/openequivariance/openequivariance/core/e3nn_lite.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"074719fb4fa857ab"}},{"code_sha256_prefix":"0367bc28e8c59a28","entry":"scatter_add_wrapper","repo":"passionlab/openequivariance","repo_kind":"official","path":"openequivariance/openequivariance/core/ConvolutionBase.py","file_url":"https://github.com/passionlab/openequivariance/blob/HEAD/openequivariance/openequivariance/core/ConvolutionBase.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"0367bc28e8c59a28"}},{"code_sha256_prefix":"1231a00297920f0e","entry":"sparse_outer_product_work","repo":"passionlab/openequivariance","repo_kind":"official","path":"openequivariance/openequivariance/core/utils.py","file_url":"https://github.com/passionlab/openequivariance/blob/HEAD/openequivariance/openequivariance/core/utils.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"1231a00297920f0e"}},{"code_sha256_prefix":"a287eaa5d7a7b744","entry":"wigner_3j","repo":"passionlab/openequivariance","repo_kind":"official","path":"openequivariance/openequivariance/core/e3nn_lite.py","file_url":"https://github.com/passionlab/openequivariance/blob/HEAD/openequivariance/openequivariance/core/e3nn_lite.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"a287eaa5d7a7b744"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}