mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Update CUDA custom op unit tests to account for recent ORT change (#6971)
This commit is contained in:
Hariharan Seshadri
GitHub
parent
694389a85d
commit
12b5ab3bab
8 changed files with 471 additions and 392 deletions
|
|
@ -319,7 +319,9 @@ set (onnxruntime_shared_lib_test_SRC
|
|||
${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_model_loading.cc
|
||||
${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_ort_format_models.cc
|
||||
${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/utils.h
|
||||
${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/utils.cc)
|
||||
${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/utils.cc
|
||||
${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/custom_op_utils.h
|
||||
${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/custom_op_utils.cc)
|
||||
|
||||
if (NOT onnxruntime_MINIMAL_BUILD)
|
||||
list(APPEND onnxruntime_shared_lib_test_SRC ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_inference.cc)
|
||||
|
|
|
|||
|
|
@ -16,8 +16,8 @@ __global__ void cuda_add_impl(int64_t N, T3* O, const T1* X, const T2* Y) {
|
|||
}
|
||||
|
||||
template <typename T1, typename T2, typename T3>
|
||||
void cuda_add(int64_t N, T3* O, const T1* X, const T2* Y) {
|
||||
cuda_add_impl<<<1, 256, 0, 0>>>(N, O, X, Y);
|
||||
void cuda_add(int64_t N, T3* O, const T1* X, const T2* Y, cudaStream_t compute_stream) {
|
||||
cuda_add_impl<<<1, 256, 0, compute_stream>>>(N, O, X, Y);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
|
|
@ -29,12 +29,12 @@ __global__ void cuda_slice_impl(const T* X, int64_t from, int64_t to, T* Y) {
|
|||
}
|
||||
|
||||
template <typename T>
|
||||
void cuda_slice(const T* X, int64_t from, int64_t to, T* Y) {
|
||||
cuda_slice_impl<T><<<1, 256, 0, 0>>>(X, from, to, Y);
|
||||
void cuda_slice(const T* X, int64_t from, int64_t to, T* Y, cudaStream_t compute_stream) {
|
||||
cuda_slice_impl<T><<<1, 256, 0, compute_stream>>>(X, from, to, Y);
|
||||
}
|
||||
|
||||
template void cuda_slice(const float*, int64_t, int64_t, float*);
|
||||
template void cuda_slice(const double*, int64_t, int64_t, double*);
|
||||
template void cuda_slice(const float*, int64_t, int64_t, float*, cudaStream_t compute_stream);
|
||||
template void cuda_slice(const double*, int64_t, int64_t, double*, cudaStream_t compute_stream);
|
||||
|
||||
template void cuda_add(int64_t, float*, const float*, const float*);
|
||||
template void cuda_add(int64_t, float*, const float*, const double*);
|
||||
template void cuda_add(int64_t, float*, const float*, const float*, cudaStream_t compute_stream);
|
||||
template void cuda_add(int64_t, float*, const float*, const double*, cudaStream_t compute_stream);
|
||||
|
|
|
|||
184
onnxruntime/test/shared_lib/custom_op_utils.cc
Normal file
184
onnxruntime/test/shared_lib/custom_op_utils.cc
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "custom_op_utils.h"
|
||||
#include "core/common/common.h"
|
||||
|
||||
#ifdef USE_CUDA
|
||||
#include <cuda_runtime.h>
|
||||
template <typename T1, typename T2, typename T3>
|
||||
void cuda_add(int64_t, T3*, const T1*, const T2*, cudaStream_t compute_stream);
|
||||
|
||||
template <typename T>
|
||||
void cuda_slice(const T*, int64_t, int64_t, T*, cudaStream_t compute_stream);
|
||||
#endif
|
||||
|
||||
void MyCustomKernel::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_Y = ort_.KernelContext_GetInput(context, 1);
|
||||
const float* X = ort_.GetTensorData<float>(input_X);
|
||||
const float* Y = ort_.GetTensorData<float>(input_Y);
|
||||
|
||||
// Setup output
|
||||
OrtTensorDimensions dimensions(ort_, input_X);
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions.data(), dimensions.size());
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output);
|
||||
int64_t size = ort_.GetTensorShapeElementCount(output_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(output_info);
|
||||
|
||||
// Do computation
|
||||
#ifdef USE_CUDA
|
||||
// Launch on stream 0 or user provided stream
|
||||
cuda_add(size, out, X, Y, compute_stream_ == nullptr ? 0 : reinterpret_cast<cudaStream_t>(compute_stream_));
|
||||
// If everything is setup correctly, custom op implementations need not have such synchronization logic.
|
||||
// To make sure custom ops and ORT CUDA kernels are implicitly synchronized, create your session with a compute stream
|
||||
// passed in via SessionOptions and use the same compute stream ti launch the custom op (as shown in this example)
|
||||
// cudaStreamSynchronize(nullptr);
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(compute_stream_);
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = X[i] + Y[i];
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
void MyCustomKernelMultipleDynamicInputs::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_Y = ort_.KernelContext_GetInput(context, 1);
|
||||
// Even though this kernel backs an operator where-in both inputs can be any type and need not be homogeneous
|
||||
// as a proof-of-concept, support the case where-in the first input is of float type and the second input
|
||||
// is of double type. Users need to extend this logic to handle any arbitrary type should the need arise.
|
||||
const float* X = ort_.GetTensorData<float>(input_X);
|
||||
const double* Y = ort_.GetTensorData<double>(input_Y);
|
||||
|
||||
// Setup output
|
||||
OrtTensorDimensions dimensions(ort_, input_X);
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions.data(), dimensions.size());
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output);
|
||||
int64_t size = ort_.GetTensorShapeElementCount(output_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(output_info);
|
||||
|
||||
// Do computation
|
||||
#ifdef USE_CUDA
|
||||
// Launch on stream 0 or user provided stream
|
||||
cuda_add(size, out, X, Y, compute_stream_ == nullptr ? 0 : reinterpret_cast<cudaStream_t>(compute_stream_));
|
||||
// If everything is setup correctly, custom op implementations need not have such synchronization logic.
|
||||
// To make sure custom ops and ORT CUDA kernels are implicitly synchronized, create your session with a compute stream
|
||||
// passed in via SessionOptions and use the same compute stream ti launch the custom op (as shown in this example)
|
||||
// cudaStreamSynchronize(nullptr);
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(compute_stream_);
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = static_cast<float>(X[i] + Y[i]);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
void MyCustomKernelWithOptionalInput::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X1 = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_X2 = ort_.KernelContext_GetInput(context, 1);
|
||||
const OrtValue* input_X3 = ort_.KernelContext_GetInput(context, 2);
|
||||
|
||||
const float* X1 = ort_.GetTensorData<float>(input_X1);
|
||||
// The second input may or may not be present
|
||||
const float* X2 = (input_X2 != nullptr) ? ort_.GetTensorData<float>(input_X2) : nullptr;
|
||||
const float* X3 = ort_.GetTensorData<float>(input_X3);
|
||||
|
||||
// Setup output
|
||||
int64_t output_dim_value = 1;
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, &output_dim_value, 1);
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
// Only CPU EP is supported in this kernel
|
||||
for (int64_t i = 0; i < output_dim_value; i++) {
|
||||
out[i] = X1[i] + (X2 != nullptr ? X2[i] : 0) + X3[i];
|
||||
}
|
||||
}
|
||||
|
||||
void MyCustomKernelWithAttributes::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const float* X = ort_.GetTensorData<float>(input_X);
|
||||
|
||||
// Setup output
|
||||
OrtTensorDimensions dimensions(ort_, input_X);
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions.data(), dimensions.size());
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output);
|
||||
int64_t size = ort_.GetTensorShapeElementCount(output_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(output_info);
|
||||
|
||||
// This kernel only supports CPU EP
|
||||
if (string_arr_ == "add") { // Test that the string attribute parsing went correctly
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = X[i] +
|
||||
float_attr_ + static_cast<float>(int_attr_) +
|
||||
floats_attr_[0] + floats_attr_[1] +
|
||||
static_cast<float>(ints_attr_[0]) + static_cast<float>(ints_attr_[1]);
|
||||
}
|
||||
} else { // if the string attribute parsing had not gone correctly - it will trigger this path and fail the test due to result mis-match
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = 0.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void custom_slice(const T* X, int64_t from, int64_t to, T* Y, void* compute_stream) {
|
||||
#ifdef USE_CUDA
|
||||
// Launch on stream 0 or user provided stream
|
||||
cuda_slice(X, from, to, Y, compute_stream == nullptr ? 0 : reinterpret_cast<cudaStream_t>(compute_stream));
|
||||
// If everything is setup correctly, custom op implementations need not have such synchronization logic.
|
||||
// To make sure custom ops and ORT CUDA kernels are implicitly synchronized, create your session with a compute stream
|
||||
// passed in via SessionOptions and use the same compute stream ti launch the custom op (as shown in this example)
|
||||
// cudaStreamSynchronize(nullptr);
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(compute_stream);
|
||||
for (auto i = from; i < to; i++) {
|
||||
Y[i - from] = X[i];
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
void SliceCustomOpKernel::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs and outputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_from = ort_.KernelContext_GetInput(context, 1);
|
||||
const OrtValue* input_to = ort_.KernelContext_GetInput(context, 2);
|
||||
OrtTensorTypeAndShapeInfo* input_X_info = ort_.GetTensorTypeAndShape(input_X);
|
||||
ONNXTensorElementDataType input_X_type = ort_.GetTensorElementType(input_X_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(input_X_info);
|
||||
#if USE_CUDA
|
||||
int64_t slice_from = 0;
|
||||
int64_t slice_to = 0;
|
||||
cudaMemcpy(&slice_from, ort_.GetTensorData<int64_t>(input_from), sizeof(int64_t), cudaMemcpyDeviceToHost);
|
||||
cudaMemcpy(&slice_to, ort_.GetTensorData<int64_t>(input_to), sizeof(int64_t), cudaMemcpyDeviceToHost);
|
||||
#else
|
||||
int64_t slice_from = *ort_.GetTensorData<int64_t>(input_from);
|
||||
int64_t slice_to = *ort_.GetTensorData<int64_t>(input_to);
|
||||
#endif
|
||||
std::vector<int64_t> output_dims = {slice_to - slice_from};
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, output_dims.data(), output_dims.size());
|
||||
// do slice
|
||||
switch (input_X_type) {
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
|
||||
|
||||
custom_slice(ort_.GetTensorData<float>(input_X), slice_from, slice_to,
|
||||
ort_.GetTensorMutableData<float>(output), compute_stream_);
|
||||
break;
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
|
||||
custom_slice(ort_.GetTensorData<double>(input_X), slice_from, slice_to,
|
||||
ort_.GetTensorMutableData<double>(output), compute_stream_);
|
||||
break;
|
||||
default:
|
||||
ORT_THROW("Unsupported input type");
|
||||
}
|
||||
}
|
||||
211
onnxruntime/test/shared_lib/custom_op_utils.h
Normal file
211
onnxruntime/test/shared_lib/custom_op_utils.h
Normal file
|
|
@ -0,0 +1,211 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/session/onnxruntime_cxx_api.h"
|
||||
#include <vector>
|
||||
|
||||
#ifdef USE_CUDA
|
||||
#include <cuda_runtime.h>
|
||||
#endif
|
||||
|
||||
struct Input {
|
||||
const char* name = nullptr;
|
||||
std::vector<int64_t> dims;
|
||||
std::vector<float> values;
|
||||
};
|
||||
|
||||
struct OrtTensorDimensions : std::vector<int64_t> {
|
||||
OrtTensorDimensions(Ort::CustomOpApi ort, const OrtValue* value) {
|
||||
OrtTensorTypeAndShapeInfo* info = ort.GetTensorTypeAndShape(value);
|
||||
std::vector<int64_t>::operator=(ort.GetTensorShape(info));
|
||||
ort.ReleaseTensorTypeAndShapeInfo(info);
|
||||
}
|
||||
};
|
||||
|
||||
struct MyCustomKernel {
|
||||
MyCustomKernel(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/, void* compute_stream)
|
||||
: ort_(ort), compute_stream_(compute_stream) {
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
void* compute_stream_;
|
||||
};
|
||||
|
||||
struct MyCustomOp : Ort::CustomOpBase<MyCustomOp, MyCustomKernel> {
|
||||
explicit MyCustomOp(const char* provider, void* compute_stream) : provider_(provider), compute_stream_(compute_stream) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernel(api, info, compute_stream_); };
|
||||
const char* GetName() const { return "Foo"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 2; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
|
||||
// Both the inputs need to be necessarily of float type
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
|
||||
};
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
void* compute_stream_;
|
||||
};
|
||||
|
||||
struct MyCustomKernelMultipleDynamicInputs {
|
||||
MyCustomKernelMultipleDynamicInputs(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/, void* compute_stream)
|
||||
: ort_(ort), compute_stream_(compute_stream) {
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
void* compute_stream_;
|
||||
};
|
||||
|
||||
struct MyCustomOpMultipleDynamicInputs : Ort::CustomOpBase<MyCustomOpMultipleDynamicInputs, MyCustomKernelMultipleDynamicInputs> {
|
||||
explicit MyCustomOpMultipleDynamicInputs(const char* provider, void* compute_stream) : provider_(provider), compute_stream_(compute_stream) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernelMultipleDynamicInputs(api, info, compute_stream_); };
|
||||
const char* GetName() const { return "Foo"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 2; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
|
||||
// Both the inputs are dynamic typed (i.e.) they can be any type and need not be
|
||||
// homogeneous
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
|
||||
};
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
void* compute_stream_;
|
||||
};
|
||||
|
||||
struct MyCustomKernelWithOptionalInput {
|
||||
MyCustomKernelWithOptionalInput(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/) : ort_(ort) {
|
||||
}
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
};
|
||||
|
||||
struct MyCustomOpWithOptionalInput : Ort::CustomOpBase<MyCustomOpWithOptionalInput, MyCustomKernelWithOptionalInput> {
|
||||
explicit MyCustomOpWithOptionalInput(const char* provider) : provider_(provider) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernelWithOptionalInput(api, info); };
|
||||
const char* GetName() const { return "FooBar"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 3; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t index) const {
|
||||
// The second input (index == 1) is optional
|
||||
if (index == 1)
|
||||
return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL;
|
||||
|
||||
return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED;
|
||||
}
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t /*index*/) const {
|
||||
return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED;
|
||||
}
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
};
|
||||
|
||||
struct MyCustomKernelWithAttributes {
|
||||
MyCustomKernelWithAttributes(Ort::CustomOpApi ort, const OrtKernelInfo* info) : ort_(ort) {
|
||||
int_attr_ = ort_.KernelInfoGetAttribute<int64_t>(info, "int_attr");
|
||||
float_attr_ = ort_.KernelInfoGetAttribute<float>(info, "float_attr");
|
||||
|
||||
ints_attr_ = ort_.KernelInfoGetAttribute<std::vector<int64_t>>(info, "ints_attr");
|
||||
floats_attr_ = ort_.KernelInfoGetAttribute<std::vector<float>>(info, "floats_attr");
|
||||
|
||||
string_arr_ = ort_.KernelInfoGetAttribute<std::string>(info, "string_attr");
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
|
||||
int64_t int_attr_;
|
||||
float float_attr_;
|
||||
|
||||
std::vector<int64_t> ints_attr_;
|
||||
std::vector<float> floats_attr_;
|
||||
|
||||
std::string string_arr_;
|
||||
};
|
||||
|
||||
struct MyCustomOpWithAttributes : Ort::CustomOpBase<MyCustomOpWithAttributes, MyCustomKernelWithAttributes> {
|
||||
explicit MyCustomOpWithAttributes(const char* provider) : provider_(provider) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernelWithAttributes(api, info); };
|
||||
const char* GetName() const { return "FooBar_Attr"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
};
|
||||
|
||||
//Slice array of floats or doubles between [from, to) and save to output
|
||||
struct SliceCustomOpKernel {
|
||||
SliceCustomOpKernel(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/, void* compute_stream)
|
||||
: ort_(ort), compute_stream_(compute_stream) {
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
void* compute_stream_;
|
||||
};
|
||||
|
||||
struct SliceCustomOp : Ort::CustomOpBase<SliceCustomOp, SliceCustomOpKernel> {
|
||||
explicit SliceCustomOp(const char* provider, void* compute_stream) : provider_(provider), compute_stream_(compute_stream) {}
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const {
|
||||
return new SliceCustomOpKernel(api, info, compute_stream_);
|
||||
};
|
||||
|
||||
const char* GetName() const { return "Slice"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 3; };
|
||||
ONNXTensorElementDataType GetInputType(size_t index) const {
|
||||
if (index == 0)
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; // input array of float or double
|
||||
else if (index == 1)
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; // slice from
|
||||
|
||||
// index 2 (keep compiler happy on Linux)
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; // slice to
|
||||
};
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const {
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
|
||||
}
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
void* compute_stream_;
|
||||
};
|
||||
|
|
@ -22,6 +22,7 @@
|
|||
#include "test_allocator.h"
|
||||
#include "test_fixture.h"
|
||||
#include "utils.h"
|
||||
#include "custom_op_utils.h"
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <Windows.h>
|
||||
|
|
@ -87,14 +88,17 @@ static void TestInference(Ort::Env& env, const std::basic_string<ORTCHAR_T>& mod
|
|||
OrtCustomOpDomain* custom_op_domain_ptr,
|
||||
const char* custom_op_library_filename,
|
||||
void** library_handle = nullptr,
|
||||
bool test_session_creation_only = false) {
|
||||
bool test_session_creation_only = false,
|
||||
void* cuda_compute_stream = nullptr) {
|
||||
Ort::SessionOptions session_options;
|
||||
|
||||
if (provider_type == 1) {
|
||||
#ifdef USE_CUDA
|
||||
Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0));
|
||||
std::cout << "Running simple inference with cuda provider" << std::endl;
|
||||
auto cuda_options = CreateDefaultOrtCudaProviderOptionsWithCustomStream(cuda_compute_stream);
|
||||
session_options.AppendExecutionProvider_CUDA(cuda_options);
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(cuda_compute_stream);
|
||||
return;
|
||||
#endif
|
||||
} else if (provider_type == 2) {
|
||||
|
|
@ -241,9 +245,11 @@ TEST(CApiTest, custom_op_handler) {
|
|||
std::vector<float> expected_values_y = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f};
|
||||
|
||||
#ifdef USE_CUDA
|
||||
MyCustomOp custom_op{onnxruntime::kCudaExecutionProvider};
|
||||
cudaStream_t compute_stream = nullptr;
|
||||
cudaStreamCreateWithFlags(&compute_stream, cudaStreamNonBlocking);
|
||||
MyCustomOp custom_op{onnxruntime::kCudaExecutionProvider, compute_stream};
|
||||
#else
|
||||
MyCustomOp custom_op{onnxruntime::kCpuExecutionProvider};
|
||||
MyCustomOp custom_op{onnxruntime::kCpuExecutionProvider, nullptr};
|
||||
#endif
|
||||
|
||||
Ort::CustomOpDomain custom_op_domain("");
|
||||
|
|
@ -251,111 +257,14 @@ TEST(CApiTest, custom_op_handler) {
|
|||
|
||||
#ifdef USE_CUDA
|
||||
TestInference<float>(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 1,
|
||||
custom_op_domain, nullptr, nullptr);
|
||||
custom_op_domain, nullptr, nullptr, false, compute_stream);
|
||||
cudaStreamDestroy(compute_stream);
|
||||
#else
|
||||
TestInference<float>(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 0,
|
||||
custom_op_domain, nullptr);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void cuda_slice(const T*, int64_t, int64_t, T*);
|
||||
|
||||
template <typename T>
|
||||
void custom_slice(const T* X, int64_t from, int64_t to, T* Y) {
|
||||
#ifdef USE_CUDA
|
||||
cuda_slice(X, from, to, Y);
|
||||
#else
|
||||
for (auto i = from; i < to; i++) {
|
||||
Y[i - from] = X[i];
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
//Slice array of floats or doubles between [from, to) and save to output
|
||||
struct SliceCustomOpKernel {
|
||||
SliceCustomOpKernel(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/) : ort_(ort) {
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context) {
|
||||
// Setup inputs and outputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_from = ort_.KernelContext_GetInput(context, 1);
|
||||
const OrtValue* input_to = ort_.KernelContext_GetInput(context, 2);
|
||||
OrtTensorTypeAndShapeInfo* input_X_info = ort_.GetTensorTypeAndShape(input_X);
|
||||
ONNXTensorElementDataType input_X_type = ort_.GetTensorElementType(input_X_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(input_X_info);
|
||||
#if USE_CUDA
|
||||
int64_t slice_from = 0;
|
||||
int64_t slice_to = 0;
|
||||
cudaMemcpy(&slice_from, ort_.GetTensorData<int64_t>(input_from), sizeof(int64_t), cudaMemcpyDeviceToHost);
|
||||
cudaMemcpy(&slice_to, ort_.GetTensorData<int64_t>(input_to), sizeof(int64_t), cudaMemcpyDeviceToHost);
|
||||
#else
|
||||
int64_t slice_from = *ort_.GetTensorData<int64_t>(input_from);
|
||||
int64_t slice_to = *ort_.GetTensorData<int64_t>(input_to);
|
||||
#endif
|
||||
std::vector<int64_t> output_dims = {slice_to - slice_from};
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, output_dims.data(), output_dims.size());
|
||||
// do slice
|
||||
switch (input_X_type) {
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
|
||||
custom_slice(ort_.GetTensorData<float>(input_X), slice_from, slice_to,
|
||||
ort_.GetTensorMutableData<float>(output));
|
||||
break;
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
|
||||
custom_slice(ort_.GetTensorData<double>(input_X), slice_from, slice_to,
|
||||
ort_.GetTensorMutableData<double>(output));
|
||||
break;
|
||||
default:
|
||||
ORT_THROW("Unsupported input type: ", input_X_type);
|
||||
}
|
||||
} // Compute
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
};
|
||||
|
||||
struct SliceCustomOp : Ort::CustomOpBase<SliceCustomOp, SliceCustomOpKernel> {
|
||||
explicit SliceCustomOp(const char* provider) : provider_(provider) {}
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const {
|
||||
return new SliceCustomOpKernel(api, info);
|
||||
};
|
||||
|
||||
const char* GetName() const { return "Slice"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 3; };
|
||||
ONNXTensorElementDataType GetInputType(size_t index) const {
|
||||
switch (index) {
|
||||
case 0:
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; // input array of float or double
|
||||
break;
|
||||
case 1:
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; // slice from
|
||||
break;
|
||||
case 2:
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; // slice to
|
||||
break;
|
||||
default:
|
||||
ORT_THROW("Invalid input index: ", index);
|
||||
}
|
||||
};
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t index) const {
|
||||
switch (index) {
|
||||
case 0:
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
|
||||
break;
|
||||
default:
|
||||
ORT_THROW("Invalid output index: ", index);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
};
|
||||
|
||||
//test custom op which accepts float and double as inputs
|
||||
TEST(CApiTest, varied_input_custom_op_handler) {
|
||||
std::vector<Input> inputs(2);
|
||||
|
|
@ -369,9 +278,11 @@ TEST(CApiTest, varied_input_custom_op_handler) {
|
|||
std::vector<float> expected_values_z = {10.0f};
|
||||
|
||||
#ifdef USE_CUDA
|
||||
SliceCustomOp slice_custom_op{onnxruntime::kCudaExecutionProvider};
|
||||
cudaStream_t compute_stream = nullptr;
|
||||
cudaStreamCreateWithFlags(&compute_stream, cudaStreamNonBlocking);
|
||||
SliceCustomOp slice_custom_op{onnxruntime::kCudaExecutionProvider, compute_stream};
|
||||
#else
|
||||
SliceCustomOp slice_custom_op{onnxruntime::kCpuExecutionProvider};
|
||||
SliceCustomOp slice_custom_op{onnxruntime::kCpuExecutionProvider, nullptr};
|
||||
#endif
|
||||
|
||||
Ort::CustomOpDomain custom_op_domain("abc");
|
||||
|
|
@ -379,7 +290,8 @@ TEST(CApiTest, varied_input_custom_op_handler) {
|
|||
|
||||
#ifdef USE_CUDA
|
||||
TestInference<float>(*ort_env, VARIED_INPUT_CUSTOM_OP_MODEL_URI, inputs, "Z",
|
||||
expected_dims_z, expected_values_z, 1, custom_op_domain, nullptr, nullptr);
|
||||
expected_dims_z, expected_values_z, 1, custom_op_domain, nullptr, nullptr, false, compute_stream);
|
||||
cudaStreamDestroy(compute_stream);
|
||||
#else
|
||||
TestInference<float>(*ort_env, VARIED_INPUT_CUSTOM_OP_MODEL_URI, inputs, "Z",
|
||||
expected_dims_z, expected_values_z, 0, custom_op_domain, nullptr);
|
||||
|
|
@ -388,9 +300,11 @@ TEST(CApiTest, varied_input_custom_op_handler) {
|
|||
|
||||
TEST(CApiTest, multiple_varied_input_custom_op_handler) {
|
||||
#ifdef USE_CUDA
|
||||
MyCustomOpMultipleDynamicInputs custom_op{onnxruntime::kCudaExecutionProvider};
|
||||
cudaStream_t compute_stream = nullptr;
|
||||
cudaStreamCreateWithFlags(&compute_stream, cudaStreamNonBlocking);
|
||||
MyCustomOpMultipleDynamicInputs custom_op{onnxruntime::kCudaExecutionProvider, compute_stream};
|
||||
#else
|
||||
MyCustomOpMultipleDynamicInputs custom_op{onnxruntime::kCpuExecutionProvider};
|
||||
MyCustomOpMultipleDynamicInputs custom_op{onnxruntime::kCpuExecutionProvider, nullptr};
|
||||
#endif
|
||||
|
||||
Ort::CustomOpDomain custom_op_domain("");
|
||||
|
|
@ -399,7 +313,8 @@ TEST(CApiTest, multiple_varied_input_custom_op_handler) {
|
|||
Ort::SessionOptions session_options;
|
||||
|
||||
#ifdef USE_CUDA
|
||||
Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0));
|
||||
auto cuda_options = CreateDefaultOrtCudaProviderOptionsWithCustomStream(compute_stream);
|
||||
session_options.AppendExecutionProvider_CUDA(cuda_options);
|
||||
#endif
|
||||
|
||||
session_options.Add(custom_op_domain);
|
||||
|
|
@ -444,6 +359,10 @@ TEST(CApiTest, multiple_varied_input_custom_op_handler) {
|
|||
for (size_t i = 0; i != total_len; ++i) {
|
||||
ASSERT_EQ(values_y[i], f[i]);
|
||||
}
|
||||
|
||||
#ifdef USE_CUDA
|
||||
cudaStreamDestroy(compute_stream);
|
||||
#endif
|
||||
}
|
||||
|
||||
TEST(CApiTest, optional_input_output_custom_op_handler) {
|
||||
|
|
@ -589,8 +508,9 @@ TEST(CApiTest, RegisterCustomOpForCPUAndCUDA) {
|
|||
std::vector<int64_t> expected_dims_y = {3, 2};
|
||||
std::vector<float> expected_values_y = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f};
|
||||
|
||||
MyCustomOp custom_op_cpu{onnxruntime::kCpuExecutionProvider};
|
||||
MyCustomOp custom_op_cuda{onnxruntime::kCudaExecutionProvider};
|
||||
MyCustomOp custom_op_cpu{onnxruntime::kCpuExecutionProvider, nullptr};
|
||||
// We are going to test session creation only - hence it is not a problem to use the default stream as the compute stream for the custom op
|
||||
MyCustomOp custom_op_cuda{onnxruntime::kCudaExecutionProvider, nullptr};
|
||||
Ort::CustomOpDomain custom_op_domain("");
|
||||
custom_op_domain.Add(&custom_op_cpu);
|
||||
custom_op_domain.Add(&custom_op_cuda);
|
||||
|
|
|
|||
|
|
@ -7,12 +7,14 @@
|
|||
// custom ops are only supported in a minimal build if explicitly enabled
|
||||
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
|
||||
|
||||
#include "core/common/common.h"
|
||||
#include "core/common/make_unique.h"
|
||||
#include "core/graph/constants.h"
|
||||
#include "core/session/onnxruntime_cxx_api.h"
|
||||
|
||||
#include "test_allocator.h"
|
||||
#include "utils.h"
|
||||
#include "custom_op_utils.h"
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
|
|
@ -21,15 +23,16 @@ extern std::unique_ptr<Ort::Env> ort_env;
|
|||
static void TestInference(Ort::Env& env, const std::basic_string<ORTCHAR_T>& model_uri,
|
||||
const std::vector<Input>& inputs, const char* output_name,
|
||||
const std::vector<int64_t>& expected_dims_y, const std::vector<float>& expected_values_y,
|
||||
Ort::CustomOpDomain& custom_op_domain) {
|
||||
Ort::CustomOpDomain& custom_op_domain, void* cuda_compute_stream = nullptr) {
|
||||
Ort::SessionOptions session_options;
|
||||
session_options.Add(custom_op_domain);
|
||||
|
||||
#ifdef USE_CUDA
|
||||
OrtCUDAProviderOptions cuda_options{};
|
||||
auto cuda_options = CreateDefaultOrtCudaProviderOptionsWithCustomStream(cuda_compute_stream);
|
||||
session_options.AppendExecutionProvider_CUDA(cuda_options);
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(cuda_compute_stream);
|
||||
#endif
|
||||
|
||||
Ort::Session session(env, model_uri.c_str(), session_options);
|
||||
|
||||
MockedOrtAllocator allocator;
|
||||
|
|
@ -72,9 +75,14 @@ TEST(OrtFormatCustomOpTests, ConvertOnnxModelToOrt) {
|
|||
const std::basic_string<ORTCHAR_T> ort_file = ORT_TSTR("testdata/foo_1.onnx.test_output.ort");
|
||||
|
||||
#ifdef USE_CUDA
|
||||
MyCustomOp custom_op{onnxruntime::kCudaExecutionProvider};
|
||||
// We need to launch our custom op in the same compute stream as the one we will be
|
||||
// passing to ORT via Session options to use for the entire session (i.e.) CUDA ORT kernels
|
||||
// will now use the same stream too
|
||||
cudaStream_t compute_stream = nullptr;
|
||||
cudaStreamCreateWithFlags(&compute_stream, cudaStreamNonBlocking);
|
||||
MyCustomOp custom_op{onnxruntime::kCudaExecutionProvider, compute_stream};
|
||||
#else
|
||||
MyCustomOp custom_op{onnxruntime::kCpuExecutionProvider};
|
||||
MyCustomOp custom_op{onnxruntime::kCpuExecutionProvider, nullptr};
|
||||
#endif
|
||||
Ort::CustomOpDomain custom_op_domain("");
|
||||
custom_op_domain.Add(&custom_op);
|
||||
|
|
@ -105,7 +113,12 @@ TEST(OrtFormatCustomOpTests, ConvertOnnxModelToOrt) {
|
|||
std::vector<int64_t> expected_dims_y = {3, 2};
|
||||
std::vector<float> expected_values_y = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f};
|
||||
|
||||
TestInference(*ort_env, ort_file, inputs, "Y", expected_dims_y, expected_values_y, custom_op_domain);
|
||||
#ifdef USE_CUDA
|
||||
TestInference(*ort_env, ort_file, inputs, "Y", expected_dims_y, expected_values_y, custom_op_domain, compute_stream);
|
||||
cudaStreamDestroy(compute_stream);
|
||||
#else
|
||||
TestInference(*ort_env, ort_file, inputs, "Y", expected_dims_y, expected_values_y, custom_op_domain, nullptr);
|
||||
#endif
|
||||
}
|
||||
#endif // if !defined(ORT_MINIMAL_BUILD)
|
||||
|
||||
|
|
@ -115,7 +128,7 @@ TEST(OrtFormatCustomOpTests, ConvertOnnxModelToOrt) {
|
|||
TEST(OrtFormatCustomOpTests, LoadOrtModel) {
|
||||
const std::basic_string<ORTCHAR_T> ort_file = ORT_TSTR("testdata/foo_1.onnx.ort");
|
||||
|
||||
MyCustomOp custom_op{onnxruntime::kCpuExecutionProvider};
|
||||
MyCustomOp custom_op{onnxruntime::kCpuExecutionProvider, nullptr};
|
||||
Ort::CustomOpDomain custom_op_domain("");
|
||||
custom_op_domain.Add(&custom_op);
|
||||
|
||||
|
|
|
|||
|
|
@ -2,117 +2,16 @@
|
|||
// Licensed under the MIT License.
|
||||
|
||||
#include "utils.h"
|
||||
#include <limits>
|
||||
|
||||
#ifdef USE_CUDA
|
||||
#include <cuda_runtime.h>
|
||||
template <typename T1, typename T2, typename T3>
|
||||
void cuda_add(int64_t, T3*, const T1*, const T2*);
|
||||
#endif
|
||||
|
||||
void MyCustomKernel::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_Y = ort_.KernelContext_GetInput(context, 1);
|
||||
const float* X = ort_.GetTensorData<float>(input_X);
|
||||
const float* Y = ort_.GetTensorData<float>(input_Y);
|
||||
|
||||
// Setup output
|
||||
OrtTensorDimensions dimensions(ort_, input_X);
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions.data(), dimensions.size());
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output);
|
||||
int64_t size = ort_.GetTensorShapeElementCount(output_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(output_info);
|
||||
|
||||
// Do computation
|
||||
#ifdef USE_CUDA
|
||||
cuda_add(size, out, X, Y);
|
||||
cudaStreamSynchronize(nullptr);
|
||||
#else
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = X[i] + Y[i];
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
void MyCustomKernelMultipleDynamicInputs::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_Y = ort_.KernelContext_GetInput(context, 1);
|
||||
// Even though this kernel backs an operator where-in both inputs can be any type and need not be homogeneous
|
||||
// as a proof-of-concept, support the case where-in the first input is of float type and the second input
|
||||
// is of double type. Users need to extend this logic to handle any arbitrary type should the need arise.
|
||||
const float* X = ort_.GetTensorData<float>(input_X);
|
||||
const double* Y = ort_.GetTensorData<double>(input_Y);
|
||||
|
||||
// Setup output
|
||||
OrtTensorDimensions dimensions(ort_, input_X);
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions.data(), dimensions.size());
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output);
|
||||
int64_t size = ort_.GetTensorShapeElementCount(output_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(output_info);
|
||||
|
||||
// Do computation
|
||||
#ifdef USE_CUDA
|
||||
cuda_add(size, out, X, Y);
|
||||
cudaStreamSynchronize(nullptr);
|
||||
#else
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = static_cast<float>(X[i] + Y[i]);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
void MyCustomKernelWithOptionalInput::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X1 = ort_.KernelContext_GetInput(context, 0);
|
||||
const OrtValue* input_X2 = ort_.KernelContext_GetInput(context, 1);
|
||||
const OrtValue* input_X3 = ort_.KernelContext_GetInput(context, 2);
|
||||
|
||||
const float* X1 = ort_.GetTensorData<float>(input_X1);
|
||||
// The second input may or may not be present
|
||||
const float* X2 = (input_X2 != nullptr) ? ort_.GetTensorData<float>(input_X2) : nullptr;
|
||||
const float* X3 = ort_.GetTensorData<float>(input_X3);
|
||||
|
||||
// Setup output
|
||||
int64_t output_dim_value = 1;
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, &output_dim_value, 1);
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
// Only CPU EP is supported in this kernel
|
||||
for (int64_t i = 0; i < output_dim_value; i++) {
|
||||
out[i] = X1[i] + (X2 != nullptr ? X2[i] : 0) + X3[i];
|
||||
}
|
||||
}
|
||||
|
||||
void MyCustomKernelWithAttributes::Compute(OrtKernelContext* context) {
|
||||
// Setup inputs
|
||||
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
|
||||
const float* X = ort_.GetTensorData<float>(input_X);
|
||||
|
||||
// Setup output
|
||||
OrtTensorDimensions dimensions(ort_, input_X);
|
||||
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions.data(), dimensions.size());
|
||||
float* out = ort_.GetTensorMutableData<float>(output);
|
||||
|
||||
OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output);
|
||||
int64_t size = ort_.GetTensorShapeElementCount(output_info);
|
||||
ort_.ReleaseTensorTypeAndShapeInfo(output_info);
|
||||
|
||||
// This kernel only supports CPU EP
|
||||
if (string_arr_ == "add") { // Test that the string attribute parsing went correctly
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = X[i] +
|
||||
float_attr_ + static_cast<float>(int_attr_) +
|
||||
floats_attr_[0] + floats_attr_[1] +
|
||||
static_cast<float>(ints_attr_[0]) + static_cast<float>(ints_attr_[1]);
|
||||
}
|
||||
} else { // if the string attribute parsing had not gone correctly - it will trigger this path and fail the test due to result mis-match
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
out[i] = 0.f;
|
||||
}
|
||||
}
|
||||
OrtCUDAProviderOptions CreateDefaultOrtCudaProviderOptionsWithCustomStream(void* cuda_compute_stream) {
|
||||
OrtCUDAProviderOptions cuda_options{
|
||||
0,
|
||||
OrtCudnnConvAlgoSearch::EXHAUSTIVE,
|
||||
std::numeric_limits<size_t>::max(),
|
||||
0,
|
||||
true,
|
||||
cuda_compute_stream != nullptr ? 1 : 0,
|
||||
cuda_compute_stream != nullptr ? cuda_compute_stream : nullptr};
|
||||
return cuda_options;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,158 +1,8 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "core/session/onnxruntime_cxx_api.h"
|
||||
#include <vector>
|
||||
|
||||
struct Input {
|
||||
const char* name = nullptr;
|
||||
std::vector<int64_t> dims;
|
||||
std::vector<float> values;
|
||||
};
|
||||
|
||||
struct OrtTensorDimensions : std::vector<int64_t> {
|
||||
OrtTensorDimensions(Ort::CustomOpApi ort, const OrtValue* value) {
|
||||
OrtTensorTypeAndShapeInfo* info = ort.GetTensorTypeAndShape(value);
|
||||
std::vector<int64_t>::operator=(ort.GetTensorShape(info));
|
||||
ort.ReleaseTensorTypeAndShapeInfo(info);
|
||||
}
|
||||
};
|
||||
|
||||
struct MyCustomKernel {
|
||||
MyCustomKernel(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/) : ort_(ort) {
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
};
|
||||
|
||||
struct MyCustomOp : Ort::CustomOpBase<MyCustomOp, MyCustomKernel> {
|
||||
explicit MyCustomOp(const char* provider) : provider_(provider) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernel(api, info); };
|
||||
const char* GetName() const { return "Foo"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 2; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
|
||||
// Both the inputs need to be necessarily of float type
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
|
||||
};
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
};
|
||||
|
||||
struct MyCustomKernelMultipleDynamicInputs {
|
||||
MyCustomKernelMultipleDynamicInputs(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/) : ort_(ort) {
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
};
|
||||
|
||||
struct MyCustomOpMultipleDynamicInputs : Ort::CustomOpBase<MyCustomOpMultipleDynamicInputs, MyCustomKernelMultipleDynamicInputs> {
|
||||
explicit MyCustomOpMultipleDynamicInputs(const char* provider) : provider_(provider) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernelMultipleDynamicInputs(api, info); };
|
||||
const char* GetName() const { return "Foo"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 2; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
|
||||
// Both the inputs are dynamic typed (i.e.) they can be any type and need not be
|
||||
// homogeneous
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
|
||||
};
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
};
|
||||
|
||||
struct MyCustomKernelWithOptionalInput {
|
||||
MyCustomKernelWithOptionalInput(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/) : ort_(ort) {
|
||||
}
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
};
|
||||
|
||||
struct MyCustomOpWithOptionalInput : Ort::CustomOpBase<MyCustomOpWithOptionalInput, MyCustomKernelWithOptionalInput> {
|
||||
explicit MyCustomOpWithOptionalInput(const char* provider) : provider_(provider) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernelMultipleDynamicInputs(api, info); };
|
||||
const char* GetName() const { return "FooBar"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 3; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t index) const {
|
||||
// The second input (index == 1) is optional
|
||||
if (index == 1)
|
||||
return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL;
|
||||
|
||||
return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED;
|
||||
}
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t /*index*/) const {
|
||||
return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED;
|
||||
}
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
};
|
||||
|
||||
struct MyCustomKernelWithAttributes {
|
||||
MyCustomKernelWithAttributes(Ort::CustomOpApi ort, const OrtKernelInfo* info) : ort_(ort) {
|
||||
int_attr_ = ort_.KernelInfoGetAttribute<int64_t>(info, "int_attr");
|
||||
float_attr_ = ort_.KernelInfoGetAttribute<float>(info, "float_attr");
|
||||
|
||||
ints_attr_ = ort_.KernelInfoGetAttribute<std::vector<int64_t>>(info, "ints_attr");
|
||||
floats_attr_ = ort_.KernelInfoGetAttribute<std::vector<float>>(info, "floats_attr");
|
||||
|
||||
string_arr_ = ort_.KernelInfoGetAttribute<std::string>(info, "string_attr");
|
||||
}
|
||||
|
||||
void Compute(OrtKernelContext* context);
|
||||
|
||||
private:
|
||||
Ort::CustomOpApi ort_;
|
||||
|
||||
int64_t int_attr_;
|
||||
float float_attr_;
|
||||
|
||||
std::vector<int64_t> ints_attr_;
|
||||
std::vector<float> floats_attr_;
|
||||
|
||||
std::string string_arr_;
|
||||
};
|
||||
|
||||
struct MyCustomOpWithAttributes : Ort::CustomOpBase<MyCustomOpWithAttributes, MyCustomKernelWithAttributes> {
|
||||
explicit MyCustomOpWithAttributes(const char* provider) : provider_(provider) {}
|
||||
|
||||
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernelWithAttributes(api, info); };
|
||||
const char* GetName() const { return "FooBar_Attr"; };
|
||||
const char* GetExecutionProviderType() const { return provider_; };
|
||||
|
||||
size_t GetInputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetInputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
private:
|
||||
const char* provider_;
|
||||
};
|
||||
OrtCUDAProviderOptions CreateDefaultOrtCudaProviderOptionsWithCustomStream(void* cuda_compute_stream = nullptr);
|
||||
|
|
|
|||
Loading…
Reference in a new issue