========================== TensorRTManager ========================== */
| 145 | |
| 146 | /* ========================== TensorRTManager ========================== */ |
| 147 | void TensorRTManager::create_trt_context( |
| 148 | mgb::CompNode cn, const TensorShapeArray& inp_shape, |
| 149 | nvinfer1::ICudaEngine* engine) { |
| 150 | bool has_no_context = (!m_context); |
| 151 | if (has_no_context) { |
| 152 | #if TENSOR_RT_MANAGE_ALL_WORKSPACE |
| 153 | m_context = {engine->createExecutionContext(), {}}; |
| 154 | #else |
| 155 | m_context = {engine->createExecutionContextWithoutDeviceMemory(), {}}; |
| 156 | #endif |
| 157 | } |
| 158 | MGB_MARK_USED_VAR(cn); |
| 159 | #if NV_TENSOR_RT_VERSION >= 6001 |
| 160 | auto profile_num = engine->getNbOptimizationProfiles(); |
| 161 | auto bindings_per_profile = engine->getNbBindings() / profile_num; |
| 162 | // choose nearest profile |
| 163 | #if NV_TENSOR_RT_VERSION >= 7200 |
| 164 | bool has_select_profile = false; |
| 165 | if (has_no_context) { |
| 166 | has_select_profile = true; |
| 167 | int profile_idx = 0; |
| 168 | if (profile_num > 1) { |
| 169 | double dist = DBL_MAX; |
| 170 | for (int i = 0; i < profile_num; i++) { |
| 171 | double d_sum = 0; |
| 172 | for (size_t j = 0; j < inp_shape.size(); ++j) { |
| 173 | double d = 0; |
| 174 | double l = 0; |
| 175 | auto min_dim = engine->getProfileDimensions( |
| 176 | j + bindings_per_profile * i, i, |
| 177 | nvinfer1::OptProfileSelector::kMIN); |
| 178 | auto max_dim = engine->getProfileDimensions( |
| 179 | j + bindings_per_profile * i, i, |
| 180 | nvinfer1::OptProfileSelector::kMAX); |
| 181 | auto opt_dim = engine->getProfileDimensions( |
| 182 | j + bindings_per_profile * i, i, |
| 183 | nvinfer1::OptProfileSelector::kOPT); |
| 184 | for (int k = 0; k < min_dim.nbDims; k++) { |
| 185 | int inp_v = static_cast<int>(inp_shape.at(j)[k]); |
| 186 | if (inp_v < min_dim.d[k] || inp_v > max_dim.d[k]) { |
| 187 | d = DBL_MAX; |
| 188 | break; |
| 189 | } else { |
| 190 | d += pow(inp_v - opt_dim.d[k], 2); |
| 191 | l += pow(opt_dim.d[k], 2); |
| 192 | } |
| 193 | } |
| 194 | if (d != DBL_MAX) { |
| 195 | d_sum += sqrt(d) / sqrt(l); |
| 196 | } else { |
| 197 | d_sum = DBL_MAX; |
| 198 | break; |
| 199 | } |
| 200 | } |
| 201 | if (d_sum < dist) { |
| 202 | profile_idx = i; |
| 203 | dist = d_sum; |
| 204 | } |
no test coverage detected