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

cuequivariance_ops/common/error_raft.hpp

@@ -0,0 +1,323 @@
+/*
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __RAFT_RT_ERROR
+#define __RAFT_RT_ERROR
+
+#pragma once
+
+#if defined(__GNUC__) && __has_include(<cxxabi.h>) && __has_include(<execinfo.h>)
+#define ENABLE_COLLECT_CALLSTACK
+#endif
+
+#include <cstdio>
+#include <iostream>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#ifdef ENABLE_COLLECT_CALLSTACK
+#include <cxxabi.h>
+#include <execinfo.h>
+
+#include <sstream>
+#endif
+
+namespace raft {
+
+/**
+ * @defgroup error_handling Exceptions & Error Handling
+ * @{
+ */
+
+/** base exception class for the whole of raft */
+class exception : public std::exception {
+ public:
+  /** default ctor */
+  explicit exception() noexcept : std::exception(), msg_() {}
+
+  /** copy ctor */
+  exception(exception const& src) noexcept : std::exception(), msg_(src.what())
+  {
+    collect_call_stack();
+  }
+
+  /** ctor from an input message */
+  explicit exception(std::string const msg) noexcept : std::exception(), msg_(std::move(msg))
+  {
+    collect_call_stack();
+  }
+
+  /** get the message associated with this exception */
+  char const* what() const noexcept override { return msg_.c_str(); }
+
+ private:
+  /** message associated with this exception */
+  std::string msg_;
+
+  /** append call stack info to this exception's message for ease of debug */
+  // Courtesy: https://www.gnu.org/software/libc/manual/html_node/Backtraces.html
+  void collect_call_stack() noexcept
+  {
+#ifdef ENABLE_COLLECT_CALLSTACK
+    constexpr int kSkipFrames    = 1;
+    constexpr int kMaxStackDepth = 64;
+    void* stack[kMaxStackDepth];  // NOLINT
+    auto depth = backtrace(stack, kMaxStackDepth);
+    std::ostringstream oss;
+    oss << std::endl << "Obtained " << (depth - kSkipFrames) << " stack frames" << std::endl;
+    char** strings = backtrace_symbols(stack, depth);
+    if (strings == nullptr) {
+      oss << "But no stack trace could be found!" << std::endl;
+      msg_ += oss.str();
+      return;
+    }
+    // Courtesy: https://panthema.net/2008/0901-stacktrace-demangled/
+    for (int i = kSkipFrames; i < depth; i++) {
+      oss << "#" << i << " in ";  // beginning of the backtrace line
+
+      char* mangled_name  = nullptr;
+      char* offset_begin  = nullptr;
+      char* offset_end    = nullptr;
+      auto backtrace_line = strings[i];
+
+      // Find parentheses and +address offset surrounding mangled name
+      // e.g. ./module(function+0x15c) [0x8048a6d]
+      for (char* p = backtrace_line; *p != 0; p++) {
+        if (*p == '(') {
+          mangled_name = p;
+        } else if (*p == '+') {
+          offset_begin = p;
+        } else if (*p == ')') {
+          offset_end = p;
+          break;
+        }
+      }
+
+      // Attempt to demangle the symbol
+      if (mangled_name != nullptr && offset_begin != nullptr && offset_end != nullptr &&
+          mangled_name + 1 < offset_begin) {
+        // Split the backtrace_line
+        *mangled_name++ = 0;
+        *offset_begin++ = 0;
+        *offset_end++   = 0;
+
+        // Demangle the name part
+        int status      = 0;
+        char* real_name = abi::__cxa_demangle(mangled_name, nullptr, nullptr, &status);
+
+        if (status == 0) {  // Success: substitute the real name
+          oss << backtrace_line << ": " << real_name << " +" << offset_begin << offset_end;
+        } else {  // Couldn't demangle
+          oss << backtrace_line << ": " << mangled_name << " +" << offset_begin << offset_end;
+        }
+        free(real_name);
+      } else {  // Couldn't match the symbol name
+        oss << backtrace_line;
+      }
+      oss << std::endl;
+    }
+    free(strings);
+    msg_ += oss.str();
+#endif
+  }
+};
+
+/**
+ * @brief Exception thrown when logical precondition is violated.
+ *
+ * This exception should not be thrown directly and is instead thrown by the
+ * RAFT_EXPECTS and  RAFT_FAIL macros.
+ *
+ */
+struct logic_error : public raft::exception {
+  explicit logic_error(char const* const message) : raft::exception(message) {}
+  explicit logic_error(std::string const& message) : raft::exception(message) {}
+};
+
+/**
+ * @brief Exception thrown when attempting to use CUDA features from a non-CUDA
+ * build
+ *
+ */
+struct non_cuda_build_error : public raft::exception {
+  explicit non_cuda_build_error(char const* const message) : raft::exception(message) {}
+  explicit non_cuda_build_error(std::string const& message) : raft::exception(message) {}
+};
+
+/**
+ * @brief Exception thrown when a CUDA error is encountered.
+ */
+struct cuda_error : public raft::exception {
+  explicit cuda_error(char const* const message) : raft::exception(message) {}
+  explicit cuda_error(std::string const& message) : raft::exception(message) {}
+};
+
+/**
+ * @}
+ */
+
+}  // namespace raft
+
+// FIXME: Need to be replaced with RAFT_FAIL
+/** macro to throw a runtime error */
+#define THROW(fmt, ...)                                                                       \
+  do {                                                                                        \
+    int size1 =                                                                               \
+      std::snprintf(nullptr, 0, "exception occurred! file=%s line=%d: ", __FILE__, __LINE__); \
+    int size2 = std::snprintf(nullptr, 0, fmt, ##__VA_ARGS__);                                \
+    if (size1 < 0 || size2 < 0)                                                               \
+      throw raft::exception("Error in snprintf, cannot handle raft exception.");              \
+    auto size = size1 + size2 + 1; /* +1 for final '\0' */                                    \
+    auto buf  = std::make_unique<char[]>(size_t(size));                                       \
+    std::snprintf(buf.get(),                                                                  \
+                  size1 + 1 /* +1 for '\0' */,                                                \
+                  "exception occurred! file=%s line=%d: ",                                    \
+                  __FILE__,                                                                   \
+                  __LINE__);                                                                  \
+    std::snprintf(buf.get() + size1, size2 + 1 /* +1 for '\0' */, fmt, ##__VA_ARGS__);        \
+    std::string msg(buf.get(), buf.get() + size - 1); /* -1 to remove final '\0' */           \
+    throw raft::exception(msg);                                                               \
+  } while (0)
+
+// FIXME: Need to be replaced with RAFT_EXPECTS
+/** macro to check for a conditional and assert on failure */
+#define ASSERT(check, fmt, ...)              \
+  do {                                       \
+    if (!(check)) THROW(fmt, ##__VA_ARGS__); \
+  } while (0)
+
+/**
+ * Macro to append error message to first argument.
+ * This should only be called in contexts where it is OK to throw exceptions!
+ */
+#define SET_ERROR_MSG(msg, location_prefix, fmt, ...)                                            \
+  do {                                                                                           \
+    int size1 = std::snprintf(nullptr, 0, "%s", location_prefix);                                \
+    int size2 = std::snprintf(nullptr, 0, "file=%s line=%d: ", __FILE__, __LINE__);              \
+    int size3 = std::snprintf(nullptr, 0, fmt, ##__VA_ARGS__);                                   \
+    if (size1 < 0 || size2 < 0 || size3 < 0)                                                     \
+      throw raft::exception("Error in snprintf, cannot handle raft exception.");                 \
+    auto size = size1 + size2 + size3 + 1; /* +1 for final '\0' */                               \
+    std::vector<char> buf(size);                                                                 \
+    std::snprintf(buf.data(), size1 + 1 /* +1 for '\0' */, "%s", location_prefix);               \
+    std::snprintf(                                                                               \
+      buf.data() + size1, size2 + 1 /* +1 for '\0' */, "file=%s line=%d: ", __FILE__, __LINE__); \
+    std::snprintf(buf.data() + size1 + size2, size3 + 1 /* +1 for '\0' */, fmt, ##__VA_ARGS__);  \
+    msg += std::string(buf.data(), buf.data() + size - 1); /* -1 to remove final '\0' */         \
+  } while (0)
+
+/**
+ * @defgroup assertion Assertion and error macros
+ * @{
+ */
+
+/**
+ * @brief Macro for checking (pre-)conditions that throws an exception when a condition is false
+ *
+ * @param[in] cond Expression that evaluates to true or false
+ * @param[in] fmt String literal description of the reason that cond is expected to be true with
+ * optional format tagas
+ * @throw raft::logic_error if the condition evaluates to false.
+ */
+#define RAFT_EXPECTS(cond, fmt, ...)                              \
+  do {                                                            \
+    if (!(cond)) {                                                \
+      std::string msg{};                                          \
+      SET_ERROR_MSG(msg, "RAFT failure at ", fmt, ##__VA_ARGS__); \
+      throw raft::logic_error(msg);                               \
+    }                                                             \
+  } while (0)
+
+/**
+ * @brief Indicates that an erroneous code path has been taken.
+ *
+ * @param[in] fmt String literal description of the reason that this code path is erroneous with
+ * optional format tagas
+ * @throw always throws raft::logic_error
+ */
+#define RAFT_FAIL(fmt, ...)                                     \
+  do {                                                          \
+    std::string msg{};                                          \
+    SET_ERROR_MSG(msg, "RAFT failure at ", fmt, ##__VA_ARGS__); \
+    throw raft::logic_error(msg);                               \
+  } while (0)
+
+/**
+ * @}
+ */
+
+#endif
+
+/**
+ * @brief Error checking macro for CUDA runtime API functions.
+ *
+ * Invokes a CUDA runtime API function call, if the call does not return
+ * cudaSuccess, invokes cudaGetLastError() to clear the error and throws an
+ * exception detailing the CUDA error that occurred
+ *
+ */
+#define RAFT_CUDA_TRY(call)                        \
+  do {                                             \
+    cudaError_t const status = call;               \
+    if (status != cudaSuccess) {                   \
+      cudaGetLastError();                          \
+      std::string msg{};                           \
+      SET_ERROR_MSG(msg,                           \
+                    "CUDA error encountered at: ", \
+                    "call='%s', Reason=%s:%s",     \
+                    #call,                         \
+                    cudaGetErrorName(status),      \
+                    cudaGetErrorString(status));   \
+      throw raft::cuda_error(msg);                 \
+    }                                              \
+  } while (0)
+
+/**
+ * @brief Debug macro to check for CUDA errors
+ *
+ * In a non-release build, this macro will synchronize the specified stream
+ * before error checking. In both release and non-release builds, this macro
+ * checks for any pending CUDA errors from previous calls. If an error is
+ * reported, an exception is thrown detailing the CUDA error that occurred.
+ *
+ * The intent of this macro is to provide a mechanism for synchronous and
+ * deterministic execution for debugging asynchronous CUDA execution. It should
+ * be used after any asynchronous CUDA call, e.g., cudaMemcpyAsync, or an
+ * asynchronous kernel launch.
+ */
+#ifndef NDEBUG
+#define RAFT_CHECK_CUDA(stream) RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
+#else
+#define RAFT_CHECK_CUDA(stream) RAFT_CUDA_TRY(cudaPeekAtLastError());
+#endif
+
+// /**
+//  * @brief check for cuda runtime API errors but log error instead of raising
+//  *        exception.
+//  */
+#define RAFT_CUDA_TRY_NO_THROW(call)                               \
+  do {                                                             \
+    cudaError_t const status = call;                               \
+    if (cudaSuccess != status) {                                   \
+      printf("CUDA call='%s' at file=%s line=%d failed with %s\n", \
+             #call,                                                \
+             __FILE__,                                             \
+             __LINE__,                                             \
+             cudaGetErrorString(status));                          \
+    }                                                              \
+  } while (0)

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