Replace GEMM backend: cublas.gemm -> cublaslt.matmul

cpp/include/raft/core/resource/cublaslt_handle.hpp

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) 2023, 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.
+ */
+#pragma once
+
+#include <cublasLt.h>
+#include <raft/core/cublas_macros.hpp>
+#include <raft/core/resource/resource_types.hpp>
+#include <raft/core/resources.hpp>
+
+#include <memory>
+
+namespace raft::resource {
+
+class cublaslt_resource : public resource {
+ public:
+  cublaslt_resource() { RAFT_CUBLAS_TRY(cublasLtCreate(&handle_)); }
+  ~cublaslt_resource() noexcept override { RAFT_CUBLAS_TRY_NO_THROW(cublasLtDestroy(handle_)); }
+  auto get_resource() -> void* override { return &handle_; }
+
+ private:
+  cublasLtHandle_t handle_;
+};
+
+/** Factory that knows how to construct a specific raft::resource to populate the res_t. */
+class cublaslt_resource_factory : public resource_factory {
+ public:
+  auto get_resource_type() -> resource_type override { return resource_type::CUBLASLT_HANDLE; }
+  auto make_resource() -> resource* override { return new cublaslt_resource(); }
+};
+
+/**
+ * @defgroup resource_cublas cuBLAS handle resource functions
+ * @{
+ */
+
+/**
+ * Load a cublasLt res_t from raft res if it exists, otherwise
+ * add it and return it.
+ * @param[in] res the raft resources object
+ * @return cublasLt handle
+ */
+inline auto get_cublaslt_handle(resources const& res) -> cublasLtHandle_t
+{
+  if (!res.has_resource_factory(resource_type::CUBLASLT_HANDLE)) {
+    res.add_resource_factory(std::make_shared<cublaslt_resource_factory>());
+  }
+  auto ret = *res.get_resource<cublasLtHandle_t>(resource_type::CUBLASLT_HANDLE);
+  return ret;
+};
+
+/**
+ * @}
+ */
+
+}  // namespace raft::resource

cpp/include/raft/core/resource/resource_types.hpp

@@ -41,6 +41,8 @@ enum resource_type {
   DEVICE_ID,               // cuda device id
   THRUST_POLICY,           // thrust execution policy
   WORKSPACE_RESOURCE,      // rmm device memory resource
+  CUBLASLT_HANDLE,         // cublasLt handle
+  USER_DEFINED,            // user-defined default-constructible resource
 
   LAST_KEY                 // reserved for the last key
 };

cpp/include/raft/core/resource/user_resource.hpp

@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2023, 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.
+ */
+#pragma once
+
+#include <raft/core/resource/resource_types.hpp>
+#include <raft/core/resources.hpp>
+
+#include <memory>
+#include <typeindex>
+
+namespace raft::resource {
+
+class user_resource : public resource {
+ public:
+  user_resource()                    = default;
+  ~user_resource() noexcept override = default;
+  auto get_resource() -> void* override { return this; }
+
+  template <typename Store>
+  auto load() -> Store*
+  {
+    std::lock_guard<std::mutex> _(lock_);
+    auto key = std::type_index{typeid(Store)};
+    auto pos = map_.find(key);
+    if (pos != map_.end()) { return reinterpret_cast<Store*>(pos->second.get()); }
+    auto store_ptr = new Store{};
+    map_[key] =
+      std::shared_ptr<void>(store_ptr, [](void* ptr) { delete reinterpret_cast<Store*>(ptr); });
+    return store_ptr;
+  }
+
+ private:
+  std::unordered_map<std::type_index, std::shared_ptr<void>> map_{};
+  std::mutex lock_{};
+};
+
+/** Factory that knows how to construct a specific raft::resource to populate the res_t. */
+class user_resource_factory : public resource_factory {
+ public:
+  auto get_resource_type() -> resource_type override { return resource_type::USER_DEFINED; }
+  auto make_resource() -> resource* override { return new user_resource(); }
+};
+
+/**
+ * Get the user-defined default-constructible resource if it exists, create it otherwise.
+ * @param[in] res the raft resources object
+ * @return a pointer to the user-defined resource.
+ */
+template <typename Store>
+auto get_user_resource(resources const& res) -> Store*
+{
+  if (!res.has_resource_factory(resource_type::USER_DEFINED)) {
+    res.add_resource_factory(std::make_shared<user_resource_factory>());
+  }
+  return res.get_resource<user_resource>(resource_type::USER_DEFINED)->load<Store>();
+};
+
+}  // namespace raft::resource