[Performance] Regarding the additional GPU memory space application in onnxtuntime for the custom CUDA operator deform conv2d.

[1193700079]

Describe the issue

The deform_conv function is from the path under mmdeploy
csrc/mmdeploy/backend_ops/tensorrt/deform_conv/trt_deform_conv_kernel.cu
The code is as follows:
trt_deform_conv_kernel.cu

It involves the allocation of GPU memory for workspace, and I'm not very clear about how to use the workspace, so I simply use the byte size of an output tensor for storage.
I feel that this is not efficient enough, because tensorRT has corresponding APIs to apply for workspace memory. I want to know if there is a better way in onnxruntime, and I hope we can discuss it together! Please advise!

To reproduce

This code is written inside the Compute function.

float *workSpace = nullptr;
CHECK(cudaMalloc(&workSpace, output_array_bytes));

the complete code is as follows:

#include <regex>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <filesystem>
#include <ctime>
#include <onnxruntime_cxx_api.h>
#include "check.h"
#include "trt_deform_conv_kernel.cu"
__constant__ int64_t g_bias;
__constant__ int64_t g_deform_groups;
__constant__ int64_t g_groups;
__constant__ int64_t g_im2col_step;
__constant__ int64_t g_N;
__constant__ int64_t g_C;
__constant__ int64_t g_inp_H;
__constant__ int64_t g_inp_W;
__constant__ int64_t g_channels_out;

#include <vector>
#include <numeric>
int64_t vectorProduct(const std::vector<int64_t> &vec)
{
    return std::accumulate(vec.begin(), vec.end(), 1LL, std::multiplies<int64_t>());
}

struct DeformConv2dKernel
{
    DeformConv2dKernel(const OrtApi &api, const OrtKernelInfo *info);

    void Compute(OrtKernelContext *context);

    Ort::CustomOpApi ort_;
    const OrtKernelInfo *info_;
    Ort::AllocatorWithDefaultOptions allocator_;

    int64_t bias_{};
    int64_t deform_groups_{};
    std::vector<int64_t> dilation_{};
    int64_t groups_{};
    int64_t im2col_step_{};
    std::vector<int64_t> padding_{};
    std::vector<int64_t> stride_{};
};

struct DeformConv2dOp : Ort::CustomOpBase<DeformConv2dOp, DeformConv2dKernel>
{
    void *CreateKernel(const OrtApi &api, const OrtKernelInfo *info) const
    {
        return new DeformConv2dKernel(api, info);
    };

    const char *GetName() const
    {
        return "MMCVDeformConv2d";
    };

    size_t GetInputTypeCount() const
    {
        return 3;
    };
    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;
    };

    const char *GetExecutionProviderType() const { return "CUDAExecutionProvider"; };
};

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);
    }
};

DeformConv2dKernel::DeformConv2dKernel(const OrtApi &api, const OrtKernelInfo *info)
    : ort_(api), info_(info)
{
    bias_ = ort_.KernelInfoGetAttribute<int64_t>(info, "bias");
    deform_groups_ = ort_.KernelInfoGetAttribute<int64_t>(info, "deform_groups");
    dilation_ = ort_.KernelInfoGetAttribute<std::vector<int64_t>>(info, "dilation");
    groups_ = ort_.KernelInfoGetAttribute<int64_t>(info, "groups");
    im2col_step_ = ort_.KernelInfoGetAttribute<int64_t>(info, "im2col_step");
    padding_ = ort_.KernelInfoGetAttribute<std::vector<int64_t>>(info, "padding");
    stride_ = ort_.KernelInfoGetAttribute<std::vector<int64_t>>(info, "stride");
    allocator_ = Ort::AllocatorWithDefaultOptions();
}

void DeformConv2dKernel::Compute(OrtKernelContext *context)
{
    const int64_t bias_ = this->bias_;
    const int64_t deform_groups_ = this->deform_groups_;
    const std::vector<int64_t> dilation_ = this->dilation_;
    const int64_t groups_ = this->groups_;
    const int64_t im2col_step_ = this->im2col_step_;
    const std::vector<int64_t> padding_ = this->padding_;
    const std::vector<int64_t> stride_ = this->stride_;

    const OrtValue *input = ort_.KernelContext_GetInput(context, 0);
    const float *input_data = reinterpret_cast<const float *>(ort_.GetTensorData<float>(input)); // 1*3*512*512

    const OrtValue *input1 = ort_.KernelContext_GetInput(context, 1);
    const float *conv_output = reinterpret_cast<const float *>(ort_.GetTensorData<float>(input1)); // 1*18*510*510

    const OrtValue *input2 = ort_.KernelContext_GetInput(context, 2);
    const float *weight = reinterpret_cast<const float *>(ort_.GetTensorData<float>(input2)); // 64*3*3*3

    OrtTensorDimensions input_dims(ort_, input);
    auto N = input_dims[0];
    auto C = input_dims[1];     
    auto inp_H = input_dims[2];
    auto inp_W = input_dims[3];

    OrtTensorDimensions conv_output_dims(ort_, input1);
    auto conv_output_N = conv_output_dims[0];
    auto conv_output_C = conv_output_dims[1];
    auto conv_output_inp_H = conv_output_dims[2];
    auto conv_output_inp_W = conv_output_dims[3];

    OrtTensorDimensions weight_dims(ort_, input2);
    auto weight_N = weight_dims[0];
    auto weight_C = weight_dims[1];
    auto kernel_h = weight_dims[2];
    auto kernel_w = weight_dims[3];

    std::vector<int64_t> output_dims = {1, 64, 510, 510};
    OrtValue *output =
        ort_.KernelContext_GetOutput(context, 0, output_dims.data(), output_dims.size());
    OrtTensorDimensions output_dims_(ort_, output);
    auto channels_out = output_dims_[1];

    float *out_ptr = ort_.GetTensorMutableData<float>(output);
    float *workSpace = nullptr;

    auto input_array_bytes = sizeof(float) * N * C * inp_H * inp_W;
    auto conv_output_array_bytes = sizeof(float) * conv_output_N * conv_output_C * conv_output_inp_H * conv_output_inp_W;
    auto weight_array_bytes = sizeof(float) * weight_N * weight_C * kernel_h * kernel_w;
    auto output_array_bytes = sizeof(float) * vectorProduct(output_dims);
    std::cout << "=========================2" << std::endl;
    CHECK(cudaMalloc(&workSpace, output_array_bytes));

    std::cout << "kernel compute starting!" << std::endl;
    cudaStream_t stream = nullptr;
    cublasHandle_t m_cublas_handle;
    cublasStatus_t status = cublasCreate(&m_cublas_handle);

    deform_conv(input_data, weight, conv_output, out_ptr, workSpace, N, C, inp_H, inp_W, channels_out, kernel_w, kernel_h, stride_[0], stride_[1], padding_[0], padding_[1], dilation_[0], dilation_[1], groups_, deform_groups_, im2col_step_, m_cublas_handle, stream);
    std::cout << "kernel compute finish!" << std::endl;
    CHECK(cudaDeviceSynchronize());
    std::cout << "=========================3" << std::endl;  
}

#include <codecvt>

std::wstring charToWstring(const char *str)
{
    typedef std::codecvt_utf8<wchar_t> convert_type;
    std::wstring_convert<convert_type, wchar_t> converter;

    return converter.from_bytes(str);
}

int main(int argc, char *argv[])
{
    bool isGPU = true;
    const DeformConv2dOp my_custom_op;

    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test2");
    Ort::SessionOptions session_options;

    Ort::CustomOpDomain domain("mmcv"); // 算子域
    domain.Add(&my_custom_op);

    session_options.Add(domain);

    std::vector<std::string>
        availableProviders = Ort::GetAvailableProviders();
    auto cudaAvailable = std::find(availableProviders.begin(), availableProviders.end(), "CUDAExecutionProvider");
    OrtCUDAProviderOptions cudaOption;
    session_options.AppendExecutionProvider_CUDA(cudaOption);
    Ort::Session session{nullptr};
    try
    {
        const std::string onnxPath = "D:/projects/study_onnxruntime/addConstant_onnxruntime/models/deform_conv2d.onnx";
        std::wstring w_modelPath = charToWstring(onnxPath.c_str());
        session = Ort::Session(env, w_modelPath.c_str(), session_options);
    }
    catch (const Ort::Exception &e)
    {
        std::cout << e.what() << std::endl;
    }

    std::cout << "Model was initialized." << std::endl;
    const size_t num_input_nodes = session.GetInputCount();   //==1
    const size_t num_output_nodes = session.GetOutputCount(); //==1

    Ort::AllocatorWithDefaultOptions allocator;

    std::vector<const char *> inputNames;
    std::vector<const char *> outputNames;
    std::vector<Ort::AllocatedStringPtr> input_names_ptr;
    std::vector<Ort::AllocatedStringPtr> output_names_ptr;
    std::vector<std::vector<int64_t>> inputShapes{};
    std::vector<std::vector<int64_t>> outputShapes{};
    Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(
        OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

    for (int i = 0; i < num_input_nodes; i++)
    {
        auto input_name = session.GetInputNameAllocated(i, allocator);
        inputNames.push_back(input_name.get());
        input_names_ptr.push_back(std::move(input_name));
        Ort::TypeInfo inputTypeInfo = session.GetInputTypeInfo(i);
        std::vector<int64_t> inputTensorShape = inputTypeInfo.GetTensorTypeAndShapeInfo().GetShape();
    }
    for (int i = 0; i < num_output_nodes; i++)
    {
        auto output_name = session.GetOutputNameAllocated(i, allocator);
        outputNames.push_back(output_name.get());
        output_names_ptr.push_back(std::move(output_name));
        Ort::TypeInfo outputTypeInfo = session.GetOutputTypeInfo(i);
        std::vector<int64_t> outputTensorShape = outputTypeInfo.GetTensorTypeAndShapeInfo().GetShape();
        outputShapes.emplace_back(outputTensorShape);
    }

    std::vector<int64_t> inputTensorShape{1, 3, 512, 512};
    size_t input_tensor_size = vectorProduct(inputTensorShape);
    std::vector<float> input_tensor_values(input_tensor_size, 1);
    std::vector<Ort::Value> inputTensors;
    inputTensors.emplace_back(Ort::Value::CreateTensor<float>(
        memoryInfo,
        input_tensor_values.data(),
        input_tensor_values.size(),
        inputTensorShape.data(),
        inputTensorShape.size()));

    std::cout << "=========================1" << std::endl;

    std::vector<Ort::Value> outputTensors = session.Run(Ort::RunOptions{nullptr},
                                                        inputNames.data(),
                                                        inputTensors.data(),
                                                        inputNames.size(),
                                                        outputNames.data(),
                                                        outputNames.size());
    std::cout << "=========================4" << std::endl;
    float *boxOutput = outputTensors[0].GetTensorMutableData<float>();

    std::cout << boxOutput[0] << std::endl;
    std::cout << boxOutput[1] << std::endl;
    std::cout << boxOutput[2] << std::endl;
    std::cout << boxOutput[1 * 64 * 510 * 510 - 1] << std::endl;
    std::cout << "==============end" << std::endl;
}

‘’‘

Urgency

No response

Platform

Windows

OS Version

11

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.15.1

ONNX Runtime API

C++

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

CUDA 11.8

Model File

No response

Is this a quantized model?

No