mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-25 22:26:24 +00:00
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
105 lines
3.6 KiB
C++
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
|