Re-write allocator -> hip_allocator

Author: BenBrock

include/spblas/allocator.hpp

@@ -5,43 +5,79 @@
 #include <stdexcept>
 #include <string>
 
-namespace spblas::detail {
+namespace spblas {
 
-// throw an exception if the hipError_t is not hipSuccess.
-void throw_if_error(hipError_t error_code) {
+namespace __rocsparse {
+
+// Throw an exception if the hipError_t is not hipSuccess.
+void throw_if_error(hipError_t error_code, std::string prefix = "") {
   if (error_code == hipSuccess) {
     return;
   }
   std::string name = hipGetErrorName(error_code);
   std::string message = hipGetErrorString(error_code);
-  throw std::runtime_error(name + ":" + message);
+  throw std::runtime_error(prefix + "HIP encountered an error " + name +
+                           ": \"" + message + "\"");
 }
 
-// throw an exception if the rocsparse_status is not rocsparse_status_success.
+// Throw an exception if the rocsparse_status is not rocsparse_status_success.
 void throw_if_error(rocsparse_status error_code) {
-#define REGISTER_ROCSPARSE_ERROR(error_name)                                   \
-  if (error_code == error_name) {                                              \
-    throw std::runtime_error(#error_name);                                     \
-  }
-
   if (error_code == rocsparse_status_success) {
     return;
+  } else if (error_code == rocsparse_status_invalid_handle) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_handle\"");
+  } else if (error_code == rocsparse_status_not_implemented) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_not_implemented\"");
+  } else if (error_code == rocsparse_status_invalid_pointer) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_pointer\"");
+  } else if (error_code == rocsparse_status_invalid_size) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_size\"");
+  } else if (error_code == rocsparse_status_memory_error) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_memory_error\"");
+  } else if (error_code == rocsparse_status_internal_error) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_internal_error\"");
+  } else if (error_code == rocsparse_status_invalid_value) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_value\"");
+  } else if (error_code == rocsparse_status_arch_mismatch) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_arch_mismatch\"");
+  } else if (error_code == rocsparse_status_zero_pivot) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_zero_pivot\"");
+  } else if (error_code == rocsparse_status_not_initialized) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_not_initialized\"");
+  } else if (error_code == rocsparse_status_type_mismatch) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_type_mismatch\"");
+  } else if (error_code == rocsparse_status_type_mismatch) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_size\"");
+  } else if (error_code == rocsparse_status_invalid_size) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_size\"");
+  } else if (error_code == rocsparse_status_invalid_size) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_size\"");
+  } else if (error_code == rocsparse_status_invalid_size) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_size\"");
+  } else if (error_code == rocsparse_status_invalid_size) {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"rocsparse_status_invalid_size\"");
+  } else {
+    throw std::runtime_error(
+        "rocSPARSE encountered an error: \"unknown error\"");
   }
-
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_invalid_handle);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_not_implemented);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_invalid_pointer);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_invalid_size);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_memory_error);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_internal_error);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_invalid_value);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_arch_mismatch);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_zero_pivot);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_not_initialized);
-  REGISTER_ROCSPARSE_ERROR(rocsparse_status_type_mismatch);
-#undef REGISTER_ROCSPARSE_ERROR
-
-  throw std::runtime_error("Unknown error from rocsparse_status");
 }
 
-} // namespace spblas::detail
+} // namespace __rocsparse
+
+} // namespace spblas

include/spblas/array.hpp

@@ -0,0 +1,74 @@
+#pragma once
+
+#include "exception.hpp"
+#include <hip/hip_runtime.h>
+#include <hip/hip_runtime_api.h>
+
+namespace spblas {
+
+namespace rocsparse {
+
+template <typename T, std::size_t Alignment = 0>
+class hip_allocator {
+public:
+  using value_type = T;
+  using pointer = T*;
+  using const_pointer = const T*;
+  using reference = T&;
+  using const_reference = const T&;
+  using size_type = std::size_t;
+  using difference_type = std::ptrdiff_t;
+
+  hip_allocator() noexcept {}
+  hip_allocator(hipStream_t stream) noexcept : stream_(stream) {}
+
+  template <typename U>
+  hip_allocator(const hip_allocator<U, Alignment>& other) noexcept
+      : stream_(other.stream()) {}
+
+  hip_allocator(const hip_allocator&) = default;
+  hip_allocator& operator=(const hip_allocator&) = default;
+  ~hip_allocator() = default;
+
+  using is_always_equal = std::false_type;
+
+  pointer allocate(std::size_t size) {
+    void* ptr;
+    hipError_t error = hipMallocAsync(&ptr, size * sizeof(T), stream());
+    throw_if_failure(error);
+
+    return reinterpret_cast<T*>(ptr);
+  }
+
+  void deallocate(pointer ptr, std::size_t n = 0) {
+    if (ptr != nullptr) {
+      hipError_t error = hipFreeAsync(ptr, stream());
+      throw_if_failure(error);
+    }
+  }
+
+  bool operator==(const hip_allocator&) const = default;
+  bool operator!=(const hip_allocator&) const = default;
+
+  template <typename U>
+  struct rebind {
+    using other = hip_allocator<U, Alignment>;
+  };
+
+  hipStream_t stream() const noexcept {
+    return stream_;
+  }
+
+private:
+  void throw_if_failure(hipError_t error) {
+    if (error != hipSuccess) {
+      throw std::bad_alloc{};
+    }
+  }
+
+  hipStream_t stream_ = nullptr;
+};
+
+} // namespace rocsparse
+
+} // namespace spblas