cuequivariance-ops-cu12 0.6.0__py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

cuequivariance_ops/VERSION

@@ -0,0 +1 @@
+0.6.0

cuequivariance_ops/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+from ._version import __version__, __git_commit__
+import os
+import sys
+import ctypes
+
+PREFERRED_LOAD_FLAG = ctypes.RTLD_LOCAL
+
+
+def root_dir():
+    try:
+        import importlib.metadata
+
+        dist = importlib.metadata.distribution("cuequivariance_ops")
+        root = dist.locate_file("cuequivariance_ops")
+    except Exception:
+        # last resort, will fail with writeable install
+        root = os.path.dirname(__file__)
+    return root
+
+
+def load_library():
+    try:
+        ctypes.CDLL(
+            os.path.join(root_dir(), "lib/libcue_ops.so"), mode=PREFERRED_LOAD_FLAG
+        )
+    except Exception as e:
+        print(f"Error while loading libcue_ops.so: {e}", file=sys.stderr)
+
+
+load_library()
+
+__all__ = ["__version__", "__git_commit__", "root_dir", "load_library"]

cuequivariance_ops/_version.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+
+import importlib.resources
+
+__version__ = (
+    importlib.resources.files("cuequivariance_ops")
+    .joinpath("VERSION")
+    .read_text()
+    .strip()
+)
+__git_commit__ = "release"

cuequivariance_ops/common/common.hpp

@@ -0,0 +1,98 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <vector>
+
+namespace kernelcatcher {
+
+namespace tensor_product {
+/**
+ * @brief a wrapper struct containing informations
+ * about tensor-product paths
+ */
+template <typename MathT>
+struct __attribute__((aligned(16))) tp_info {
+  /** offsets into `path_offsets_and_dims` for each "target" */
+  const int32_t* __restrict__ path_csr_offsets{nullptr};
+  /** "sources" of all paths and their offsets and dimensions */
+  const int32_t* __restrict__ path_offsets_and_dims{nullptr};
+  /** clebsch-gordan values for each path */
+  const MathT* __restrict__ path_cg_values{nullptr};
+  /** number of "target" segments */
+  int32_t num_target_segments{0};
+  /** number of path (i.e. all paths between segments) */
+  int32_t num_paths{0};
+};  // struct tp_info
+
+enum class ConnectionModeT : uint8_t {
+  kUVW = 0,
+  // UVW with U spherical harmonic
+  k1VW,  // NOLINT
+  // UVW with V spherical harmonic
+  kU1W,  // NOLINT
+  kUVU,
+  kUVV,
+  kUUW,
+  kUUU,
+  // FullTP, no weight
+  kUVUV,
+  // FullTP, U spherical harmonic
+  k1V1V,
+  // FullTP, V spherical harmonic
+  kU1U1,
+  // Linear
+  kUUVV,
+};
+}  // namespace tensor_product
+
+namespace symmetric_tensor_contraction {
+/**
+ * @brief a wrapper struct containing informations
+ * about tensor-product paths
+ */
+template <typename DataT>
+struct __attribute__((aligned(16))) clebsch_gordan_tensor {
+  const DataT* __restrict__ cg_values{nullptr};
+  const int16_t* __restrict__ cg_indices{nullptr};
+  const int32_t* __restrict__ cg_offsets{nullptr};
+  int32_t total_output_irreps{0};
+};  // struct clebsch_gordan_tensor
+}  // namespace symmetric_tensor_contraction
+
+namespace batch_linear {
+struct __attribute__((aligned(8))) MatrixLayout {
+  int32_t size_row;  // uncontracted mode
+  int32_t size_col;  // contracted mode
+};
+
+struct __attribute__((aligned(8))) IndexOffset {
+  int32_t start;
+  int32_t end;
+};
+
+enum class GemvModeT : std::uint8_t { kUVV = 0, kUUV = 1 };
+enum class WeightSharedModeT : std::int32_t { kShared = 0, kIndexed = 1, kBatched = 2 };
+/**
+ * @brief a wrapper struct containing informations
+ * about tensor-product paths
+ */
+template <typename DataT>
+struct __attribute__((aligned(16))) batch_linear_info {
+  const MatrixLayout* __restrict__ layouts{nullptr};
+  const IndexOffset* __restrict__ index_offsets{nullptr};
+  const int32_t* __restrict__ indices{nullptr};
+  const DataT* __restrict__ alpha{nullptr};
+};  // struct batch_linear_info
+
+}  // namespace batch_linear
+}  // namespace kernelcatcher

cuequivariance_ops/common/nvtx.hpp

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+
+namespace kernelcatcher::utils {
+
+/**
+ * @brief Push a named nvtx range
+ * @param name range name
+ */
+void push_range(const char* name);
+
+/** Pop the latest range */
+void pop_range();
+
+struct range_guard {
+  range_guard(const char* name) { push_range(name); }
+  ~range_guard() { pop_range(); }
+  range_guard(range_guard const&)            = delete;
+  range_guard& operator=(range_guard const&) = delete;
+};
+
+}  // namespace kernelcatcher::utils

cuequivariance_ops/equivariance/batch_dimension.hh

@@ -0,0 +1,15 @@
+/*
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+
+namespace kernelcatcher::utils {
+
+enum class BatchDimension : int { kBatched = 0, kShared = 1, kIndexed = 2 };
+
+}  // namespace kernelcatcher::utils

cuequivariance_ops/equivariance/dtypes.hh

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+
+#include <iostream>
+
+namespace kernelcatcher::utils {
+
+enum class Datatype : int {
+  kFloat32  = 0,
+  kFloat64  = 1,
+  kFloat16  = 2,
+  kBFloat16 = 3,
+  kInt32    = 4,
+  kInt64    = 5
+};
+
+inline int size_of(Datatype dtype)
+{
+  switch (dtype) {
+    case Datatype::kFloat32: return 4;
+    case Datatype::kFloat64: return 8;
+    case Datatype::kFloat16: return 2;
+    case Datatype::kBFloat16: return 2;
+    case Datatype::kInt32: return 4;
+    case Datatype::kInt64: return 8;
+    default: return -1;
+  }
+}
+
+inline std::ostream& operator<<(std::ostream& s, Datatype const& d)
+{
+  switch (d) {
+    case Datatype::kFloat32: return s << "float";
+    case Datatype::kFloat64: return s << "double";
+    case Datatype::kFloat16: return s << "k_fp16";
+    case Datatype::kBFloat16: return s << "k_bf16";
+    case Datatype::kInt32: return s << "kc_int32";
+    case Datatype::kInt64: return s << "kc_int64";
+  }
+  return s << "unknown_datatype";
+}
+
+inline bool is_real(Datatype const& d)
+{
+  switch (d) {
+    case Datatype::kFloat32:
+    case Datatype::kFloat64:
+    case Datatype::kFloat16:
+    case Datatype::kBFloat16: return true;
+    case Datatype::kInt32:
+    case Datatype::kInt64: return false;
+  }
+  return false;  // Default case, should not be reached
+}
+
+inline bool is_integral(Datatype const& d) { return !is_real(d); }
+
+}  // namespace kernelcatcher::utils

cuequivariance_ops/equivariance/fused_tensor_product.cuh

@@ -0,0 +1,297 @@
+/*
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+
+#include "../common/common.hpp"
+
+#include <algorithm>
+#include <limits>
+
+namespace kernelcatcher::tensor_product {
+
+struct __attribute__((aligned(16))) tp_data_sizes {
+  int64_t batch_size;
+  bool shared_a;
+  bool shared_b;
+  bool shared_w;
+  int32_t stride_a;
+  int32_t stride_b;
+  int32_t stride_w;
+  int32_t stride_o;
+};  // struct tp_data_sizes
+
+template <typename DataAT, typename DataBT, typename DataWeightT, typename DataOutT, typename MathT>
+void fused_tensor_product_fwd(DataOutT* out,
+                              const DataAT* in_a,
+                              const DataBT* in_b,
+                              const DataWeightT* weight,
+                              ConnectionModeT mode,
+                              const tp_info<MathT>& info,
+                              const tp_data_sizes& sizes,
+                              cudaStream_t stream);
+
+template <typename DataAT, typename DataBT, typename DataWeightT, typename DataOutT, typename MathT>
+void fused_tensor_product_bwd(DataAT* grad_in_a,
+                              DataBT* grad_in_b,
+                              DataWeightT* grad_weight,
+                              const DataOutT* grad_out,
+                              const DataAT* in_a,
+                              const DataBT* in_b,
+                              const DataWeightT* weight,
+                              ConnectionModeT mode,
+                              const tp_info<MathT>& info_bwd_dgrad_a,
+                              const tp_info<MathT>& info_bwd_dgrad_b,
+                              const tp_info<MathT>& info_bwd_dgrad_w,
+                              const tp_data_sizes& sizes,
+                              cudaStream_t stream);
+
+template <typename DataAT, typename DataBT, typename DataWeightT, typename DataOutT, typename MathT>
+void fused_tensor_product_bwd_bwd(DataAT* grad_in_a,
+                                  DataBT* grad_in_b,
+                                  DataWeightT* grad_weight,
+                                  DataOutT* grad_grad_out,
+                                  const DataAT* grad_grad_in_a,
+                                  const DataBT* grad_grad_in_b,
+                                  const DataWeightT* grad_grad_weight,
+                                  const DataOutT* grad_out,
+                                  const DataAT* in_a,
+                                  const DataBT* in_b,
+                                  const DataWeightT* weight,
+                                  ConnectionModeT mode,
+                                  const tp_info<MathT>& info_fwd,
+                                  const tp_info<MathT>& info_bwd_dgrad_a,
+                                  const tp_info<MathT>& info_bwd_dgrad_b,
+                                  const tp_info<MathT>& info_bwd_dgrad_w,
+                                  const tp_data_sizes& sizes,
+                                  cudaStream_t stream);
+
+extern template void fused_tensor_product_bwd_bwd<float, float, float, float, float>(
+  float*,
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd_bwd<float, float, float, float, double>(
+  float*,
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd_bwd<double, double, double, double, double>(
+  double*,
+  double*,
+  double*,
+  double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd_bwd<__half, __half, __half, __half, float>(
+  __half*,
+  __half*,
+  __half*,
+  __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void
+fused_tensor_product_bwd_bwd<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16, __nv_bfloat16, float>(
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<float, float, float, float, float>(
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<float, float, float, float, double>(
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<double, double, double, double, double>(
+  double*,
+  double*,
+  double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<__half, __half, __half, __half, float>(
+  __half*,
+  __half*,
+  __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void
+fused_tensor_product_bwd<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16, __nv_bfloat16, float>(
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_fwd<float, float, float, float, float>(
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void fused_tensor_product_fwd<float, float, float, float, double>(
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void fused_tensor_product_fwd<double, double, double, double, double>(
+  double*,
+  const double*,
+  const double*,
+  const double*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_fwd<__half, __half, __half, __half, float>(
+  __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void
+fused_tensor_product_fwd<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16, __nv_bfloat16, float>(
+  __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+}  // namespace kernelcatcher::tensor_product

cuequivariance_ops/equivariance/indexed_linear.hh

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+
+#include "dtypes.hh"  // for Datatype
+
+namespace kernelcatcher::equivariance::indexed_linear {
+using namespace kernelcatcher::utils;  // for Datatype
+
+#define KC_INDEXED_LINEAR_DECL_ARGUMENTS                                                  \
+  const void *ptr_A, const void *ptr_B, const int *counts, void *ptr_C, Datatype dtype_A, \
+    Datatype dtype_B, Datatype dtype_D, int Z, int C, int u, int v, double coefficient,   \
+    Datatype math_dtype, void *workspace, size_t workspace_size, void *stream
+
+#define KC_INDEXED_LINEAR_ARGUMENTS                                                            \
+  ptr_A, ptr_B, counts, ptr_C, dtype_A, dtype_B, dtype_D, Z, C, u, v, coefficient, math_dtype, \
+    workspace, workspace_size, stream
+
+int run_indexed_linear_B(  // ptr_A  Zu
+                           // ptr_B  Cuv or Cvu
+                           // ptr_C  Zv
+  bool transpose_B,
+  KC_INDEXED_LINEAR_DECL_ARGUMENTS);
+
+int run_indexed_linear_C(  // ptr_A  Zu
+                           // ptr_B  Zv
+                           // ptr_C  Cuv
+  KC_INDEXED_LINEAR_DECL_ARGUMENTS);
+
+}  // namespace kernelcatcher::equivariance::indexed_linear