onnxruntime/onnxruntime/python/onnxruntime_pybind_state_common.cc
Dmitri Smirnov 950fe5e28b
Implement SparseTensor and infrastructure suppport and advance ONNX commit (#8038)
SparseTensor support
  Implement Builder pattern
  Fix support for 1-D and 2-D COO indices
  Implement and test CSR support.
  Handle shape inference for SparseTensors
  Implement conversion for COO, CSR and tests.
  Address the case where constant sparse initializer is the output.
  Implement test infra for SparseTensors
  Implement SparseDenseMatMul for Csr and COO and tested it.
  Add hash for SparseToDenseMatMul
  Finish shared provider refactor
  Refactor GetOrCreate to Create
  Working on py interface
  Expose OrtDevice and use it in allocate_numpy
	Adjust Sparse interfaces, add support for string SparseTensor. Add tests.
	Add and test to_cuda()
	Add accessors to format specific indices
	Test values and indices views, read-only flag, after GC access
	Add sparse related methods to OrtValue
	Re-work SparseTensor wrapper, add OrtValue methods
	Rework numpy_array_to_cuda/to_cpu
	Add run_with_ort_values
	Add models and test sparse_mat_mul with run_with_ort_values
	Refactor sparse tensor to use a single buffer
        Ifdef x86 Eigen CSR sparse matmul implementation
        Exclude broken test, check for string type when copying cross device
       Split pybind schema, regenerate docs, add exclusion
       Conditionally exclude schema module
       Update docs fix cuda build
       Add test to a filter and renerate JS docs
      Add conversion and test string support for sparse tensors
      Exclude conversion utils from minimal build
      Add CUDA Memcpy and adjust provider interfaces
2021-07-22 15:24:36 -07:00

105 lines
3.6 KiB
C++

#include "onnxruntime_pybind_exceptions.h"
#include "onnxruntime_pybind_state_common.h"
#include "core/framework/arena_extend_strategy.h"
namespace onnxruntime {
namespace python {
namespace py = pybind11;
const std::string onnxruntime::python::SessionObjectInitializer::default_logger_id = "Default";
#ifdef USE_OPENVINO
// TODO remove deprecated global config
std::string openvino_device_type;
#endif
#ifdef USE_NUPHAR
// TODO remove deprecated global config
std::string nuphar_settings;
#endif
// TODO remove deprecated global config
OrtDevice::DeviceId cuda_device_id = 0;
// TODO remove deprecated global config
size_t gpu_mem_limit = std::numeric_limits<size_t>::max();
#if defined(USE_CUDA) || defined(USE_ROCM)
#ifdef USE_CUDA
// TODO remove deprecated global config
OrtCudnnConvAlgoSearch cudnn_conv_algo_search = OrtCudnnConvAlgoSearch::EXHAUSTIVE;
// TODO remove deprecated global config
bool do_copy_in_default_stream = true;
onnxruntime::CUDAExecutionProviderExternalAllocatorInfo external_allocator_info{};
#endif
#ifdef USE_ROCM
#include "core/providers/rocm/rocm_execution_provider.h"
#include "core/providers/rocm/rocm_allocator.h"
onnxruntime::ROCMExecutionProviderExternalAllocatorInfo external_allocator_info{};
#endif
// TODO remove deprecated global config
onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo;
#endif
#ifdef ENABLE_TRAINING
static void DlpackCapsuleDestructor(PyObject* data) {
DLManagedTensor* dlmanged_tensor = reinterpret_cast<DLManagedTensor*>(
PyCapsule_GetPointer(data, "dltensor"));
if (dlmanged_tensor) {
// The dlmanged_tensor has not been consumed, call deleter ourselves.
dlmanged_tensor->deleter(const_cast<DLManagedTensor*>(dlmanged_tensor));
} else {
// The dlmanged_tensor has been consumed,
// PyCapsule_GetPointer has set an error indicator.
PyErr_Clear();
}
}
// Allocate a new Capsule object, which takes the ownership of OrtValue.
// Caller is responsible for releasing.
// This function calls OrtValueToDlpack(...).
PyObject* ToDlpack(OrtValue ort_value) {
DLManagedTensor* dlmanaged_tensor = dlpack::OrtValueToDlpack(ort_value);
return PyCapsule_New(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor);
}
// Consume a Capsule object and claims the ownership of its underlying tensor to
// create a OrtValue. This function calls DlpackToOrtValue(...) to do the conversion.
OrtValue FromDlpack(PyObject* dlpack_tensor, const bool is_bool_tensor) {
// Extract DLPack tensor pointer from the capsule carrier.
DLManagedTensor* dlmanaged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(dlpack_tensor, "dltensor");
OrtValue ort_value = dlpack::DlpackToOrtValue(dlmanaged_tensor, is_bool_tensor);
// Make sure this capsule will never be used again.
PyCapsule_SetName(dlpack_tensor, "used_dltensor");
return ort_value;
}
#endif
void PySparseTensor::Init(std::unique_ptr<SparseTensor>&& instance) {
auto sparse_tensor(std::move(instance));
auto ml_type = DataTypeImpl::GetType<SparseTensor>();
ort_value_.Init(sparse_tensor.get(), ml_type, ml_type->GetDeleteFunc());
sparse_tensor.release();
}
PySparseTensor::~PySparseTensor() {
// pybind11 will deref and potentially destroy its objects
// that we use to hold a reference and it may throw python errors
// so we want to do it in a controlled manner
auto None = py::none();
for (auto& obj : backing_storage_) {
try {
obj = None;
} catch (py::error_already_set& ex) {
// we need it mutable to properly log and discard it
ex.discard_as_unraisable(__func__);
}
}
}
} // namespace python
} // namespace onnxruntime