MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / bindAlgorithm

Function bindAlgorithm

python/src/infer/pyAlgorithmSelector.cpp:144–188  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

142}
143
144void bindAlgorithm(py::module& m)
145{
146 // IAlgorithmIOInfo
147 py::class_<IAlgorithmIOInfo, std::unique_ptr<IAlgorithmIOInfo, py::nodelete>>(
148 m, "IAlgorithmIOInfo", IAlgorithmIOInfoDOC::descr, py::module_local())
149 .def_property_readonly("tensor_format",
150 utils::deprecateMember(
151 &IAlgorithmIOInfo::getTensorFormat, "the strides, data type, and vectorization information"))
152 .def_property_readonly("dtype", &IAlgorithmIOInfo::getDataType)
153 .def_property_readonly("strides", &IAlgorithmIOInfo::getStrides)
154 .def_property_readonly("vectorized_dim", &IAlgorithmIOInfo::getVectorizedDim)
155 .def_property_readonly("components_per_element", &IAlgorithmIOInfo::getComponentsPerElement);
156
157 // IAlgorithmVariant
158 py::class_<IAlgorithmVariant, std::unique_ptr<IAlgorithmVariant, py::nodelete>>(
159 m, "IAlgorithmVariant", IAlgorithmVariantDOC::descr, py::module_local())
160 .def_property_readonly("implementation", &IAlgorithmVariant::getImplementation)
161 .def_property_readonly("tactic", &IAlgorithmVariant::getTactic);
162
163 // IAlgorithmContext
164 py::class_<IAlgorithmContext, std::unique_ptr<IAlgorithmContext, py::nodelete>>(
165 m, "IAlgorithmContext", IAlgorithmContextDoc::descr, py::module_local())
166 .def_property_readonly("name", &IAlgorithmContext::getName)
167 .def("get_shape", lambdas::get_shape, "index"_a, IAlgorithmContextDoc::get_shape)
168 .def_property_readonly("num_inputs", &IAlgorithmContext::getNbInputs)
169 .def_property_readonly("num_outputs", &IAlgorithmContext::getNbOutputs);
170
171 // IAlgorithm
172 py::class_<IAlgorithm, std::unique_ptr<IAlgorithm, py::nodelete>>(
173 m, "IAlgorithm", IAlgorithmDoc::descr, py::module_local())
174 .def("get_algorithm_io_info", &IAlgorithm::getAlgorithmIOInfoByIndex, "index"_a,
175 IAlgorithmDoc::get_algorithm_io_info, py::return_value_policy::reference_internal)
176 .def_property_readonly("algorithm_variant", &IAlgorithm::getAlgorithmVariant)
177 .def_property_readonly("timing_msec", &IAlgorithm::getTimingMSec)
178 .def_property_readonly("workspace_size", &IAlgorithm::getWorkspaceSize);
179
180 // IAlgorithmSelector
181 py::class_<IAlgorithmSelector, IAlgorithmSelectorTrampoline>(
182 m, "IAlgorithmSelector", IAlgorithmSelectorDoc::descr, py::module_local())
183 .def(py::init_alias<>()) // Always initialize trampoline class.
184 .def(
185 "select_algorithms", &select_algorithms, "context"_a, "choices"_a, IAlgorithmSelectorDoc::select_algorithms)
186 .def("report_algorithms", &report_algorithms, "contexts"_a, "choices"_a,
187 IAlgorithmSelectorDoc::report_algorithms);
188} // bindAlgorithm
189} // namespace tensorrt

Callers 1

PYBIND11_MODULEFunction · 0.85

Calls 1

deprecateMemberFunction · 0.85

Tested by

no test coverage detected