google-research

1
// Copyright 2024 The Google Research Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14

15
#if GOOGLE_CUDA
16
#include <cusparse.h>
17

18
#include <unordered_map>
19

20
#include "absl/container/flat_hash_map.h"
21
#include "sparse/ops/cc/common.h"
22

23
namespace sgk {
24

25
namespace {
26

27
using HandleMap = absl::flat_hash_map<cudaStream_t, cusparseHandle_t>;
28

29
// Get singleton map of cuda streams to cusparse handles.
30
HandleMap* GetHandleMap() {
31
  static HandleMap handle_map;
32
  return &handle_map;
33
}
34

35
}  // namespace
36

37
cusparseHandle_t GetHandleForStream(cudaStream_t stream) {
38
  HandleMap* handle_map = GetHandleMap();
39
  auto it = handle_map->find(stream);
40
  if (it == handle_map->end()) {
41
    // Allocate new cusparse handle and set the stream.
42
    cusparseHandle_t new_handle;
43
    CUSPARSE_CALL(cusparseCreate(&new_handle));
44
    CUSPARSE_CALL(cusparseSetStream(new_handle, stream));
45
    it = handle_map->insert(std::make_pair(stream, new_handle)).first;
46
  }
47
  return it->second;
48
}
49

50
}  // namespace sgk
51
#endif  // GOOGLE_CUDA
52