cuquantum_executor.cu 2.63 KB
Newer Older
1
/** ExaTN: Tensor Runtime: Tensor network executor: NVIDIA cuQuantum
2
REVISION: 2021/12/22
3
4
5
6
7
8
9
10

Copyright (C) 2018-2021 Dmitry Lyakh
Copyright (C) 2018-2021 Oak Ridge National Laboratory (UT-Battelle)

Rationale:

**/

11
#ifdef CUQUANTUM
12

13
14
15
16
17
#include <cutensornet.h>
#include <cutensor.h>
#include <cuda_runtime.h>

#include <vector>
18
#include <unordered_map>
19
20
21

#include <iostream>

22
23
#include "talshxx.hpp"

24
25
#include "cuquantum_executor.hpp"

26
27
28
29
30
31
32
#define HANDLE_CTN_ERROR(x)                                          \
{ const auto err = x;                                                 \
  if( err != CUTENSORNET_STATUS_SUCCESS )                              \
{ printf("Error: %s in line %d\n", cutensornetGetErrorString(err), __LINE__); std::abort(); } \
};


33
34
35
namespace exatn {
namespace runtime {

36
37
38
39
40
struct TensorDescriptor {
 std::vector<int32_t> modes;
 std::vector<int64_t> extents;
};

41
42
struct TensorNetworkReq {
 std::shared_ptr<numerics::TensorNetwork> network;
43
44
45
46
47
48
49
 std::unordered_map<numerics::TensorHashType,TensorDescriptor> tensor_descriptors;
 std::unordered_map<int32_t,int64_t> index_extents;
 cutensornetNetworkDescriptor_t net_descriptor;
 cutensornetContractionOptimizerConfig_t opt_config;
 cutensornetContractionOptimizerInfo_t opt_info;
 cutensornetContractionPlan_t comp_plan;
 cudaStream_t stream;
50
51
52
};


53
54
55
CuQuantumExecutor::CuQuantumExecutor()
{
 const size_t version = cutensornetGetVersion();
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
 std::cout << "#DEBUG(exatn::runtime::CuQuantumExecutor): cuTensorNet backend version " << version << std::endl;

 int num_gpus = 0;
 auto error_code = talshDeviceCount(DEV_NVIDIA_GPU,&num_gpus); assert(error_code == TALSH_SUCCESS);
 for(int i = 0; i < num_gpus; ++i){
  if(talshDeviceState(i,DEV_NVIDIA_GPU) >= DEV_ON) gpus.emplace_back(i);
 }
 std::cout << "#DEBUG(exatn::runtime::CuQuantumExecutor): Number of available GPUs = " << gpus.size() << std::endl;

 ctn_handles.resize(gpus.size());
 for(const auto & gpu_id: gpus){
  auto cuda_error = cudaSetDevice(gpu_id); assert(cuda_error == cudaSuccess);
  HANDLE_CTN_ERROR(cutensornetCreate((cutensornetHandle_t*)(&ctn_handles[gpu_id])));
 }
 std::cout << "#DEBUG(exatn::runtime::CuQuantumExecutor): Created cuTensorNet contexts for all available GPUs" << std::endl;

}


CuQuantumExecutor::~CuQuantumExecutor()
{
 bool success = sync(); assert(success);
 for(const auto & gpu_id: gpus){
  auto cuda_error = cudaSetDevice(gpu_id); assert(cuda_error == cudaSuccess);
  HANDLE_CTN_ERROR(cutensornetDestroy((cutensornetHandle_t)(ctn_handles[gpu_id])));
 }
 ctn_handles.clear();
 gpus.clear();
}


bool CuQuantumExecutor::sync()
{
 bool success = true;
 //`Finish
 return success;
92
93
94
95
96
}

} //namespace runtime
} //namespace exatn

97
#endif //CUQUANTUM