Implement ONNX QcQuantizeOp in C++

Signed-off-by: Michael Tuttle <[email protected]> Signed-off-by: Bharath Ramaswamy <[email protected]> Co-authored-by: Bharath Ramaswamy <[email protected]>
quic · Feb 17, 2023 · 3a2f2b8 · 3a2f2b8
1 parent b16e40f
commit 3a2f2b8
Show file tree

Hide file tree

Showing 13 changed files with 972 additions and 181 deletions.
diff --git a/TrainingExtensions/onnx/CMakeLists.txt b/TrainingExtensions/onnx/CMakeLists.txt
@@ -38,3 +38,73 @@ find_package(Python3 COMPONENTS Interpreter)
 
 add_subdirectory(src/python)
 add_subdirectory(test)
+
+
+add_library(OnnxCppOps SHARED
+        src/QcQuantizeOp.h
+        src/AimetOpUtils.h
+        src/AimetOpUtils.cpp
+        src/QcQuantizeInfo.cpp
+        src/QcQuantizeInfo.h
+        src/QcQuantizeOp.cpp)
+
+
+find_path(ONNXRUNTIME_INC "onnxruntime_cxx_api.h" PATH_SUFFIXES onnxruntime_headers/include)
+
+
+target_include_directories(OnnxCppOps PRIVATE
+        ${ONNXRUNTIME_INC}
+        ${CMAKE_CURRENT_SOURCE_DIR}/../common/include
+        )
+
+if (ENABLE_CUDA)
+    target_link_libraries(OnnxCppOps PUBLIC
+            MoDlQuantization
+            MoDlQuantizationCuda
+            CUDA::cublas
+            )
+else (ENABLE_CUDA)
+    target_link_libraries(OnnxCppOps PUBLIC
+            MoDlQuantization
+            )
+endif (ENABLE_CUDA)
+
+
+set_target_properties(OnnxCppOps PROPERTIES
+        OUTPUT_NAME "aimet_onnxrt_ops"
+        LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/artifacts/aimet_common"
+        )
+
+
+Python3_add_library(quant_info SHARED
+        src/QcQuantizeInfo.h
+        src/QcQuantizeInfo.cpp)
+
+
+if (ENABLE_CUDA)
+    target_link_libraries(quant_info
+            PUBLIC
+            MoDlQuantization
+            MoDlQuantizationCuda
+            CUDA::cublas
+            )
+
+else (ENABLE_CUDA)
+    target_link_libraries(quant_info
+            PUBLIC
+            MoDlQuantization
+            )
+
+endif (ENABLE_CUDA)
+
+
+set_target_properties(quant_info
+        PROPERTIES
+        OUTPUT_NAME "quant_info"
+        SUFFIX ".${Python3_SOABI}.so"
+        LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/artifacts/aimet_common/"
+        )
+
+install(TARGETS quant_info
+        LIBRARY DESTINATION ${AIMET_INSTALL_DIR}/lib/python/aimet_common
+        )
diff --git a/TrainingExtensions/onnx/src/AimetOpUtils.cpp b/TrainingExtensions/onnx/src/AimetOpUtils.cpp
@@ -0,0 +1,49 @@
+//==============================================================================
+//
+//  @@-COPYRIGHT-START-@@
+//
+//  Copyright (c) 2023, Qualcomm Innovation Center, Inc. All rights reserved.
+//
+//  Redistribution and use in source and binary forms, with or without
+//  modification, are permitted provided that the following conditions are met:
+//
+//  1. Redistributions of source code must retain the above copyright notice,
+//     this list of conditions and the following disclaimer.
+//
+//  2. Redistributions in binary form must reproduce the above copyright notice,
+//     this list of conditions and the following disclaimer in the documentation
+//     and/or other materials provided with the distribution.
+//
+//  3. Neither the name of the copyright holder nor the names of its contributors
+//     may be used to endorse or promote products derived from this software
+//     without specific prior written permission.
+//
+//  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+//  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+//  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+//  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+//  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+//  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+//  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+//  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+//  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+//  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+//  POSSIBILITY OF SUCH DAMAGE.
+//
+//  SPDX-License-Identifier: BSD-3-Clause
+//
+//  @@-COPYRIGHT-END-@@
+//
+//==============================================================================
+
+#include "AimetOpUtils.h"
+
+template <typename T>
+void copyInputTensorsToOutputTensors(const T* inTensor, size_t count, T* outTensor)
+{
+    // copy input_tensor to output_tensor
+    std::copy(inTensor, inTensor + count, outTensor);
+}
+
+
+template void copyInputTensorsToOutputTensors(const float* inTensor, size_t count, float* outTensor);
diff --git a/TrainingExtensions/onnx/src/AimetOpUtils.h b/TrainingExtensions/onnx/src/AimetOpUtils.h
@@ -0,0 +1,101 @@
+//==============================================================================
+//
+//  @@-COPYRIGHT-START-@@
+//
+//  Copyright (c) 2023, Qualcomm Innovation Center, Inc. All rights reserved.
+//
+//  Redistribution and use in source and binary forms, with or without
+//  modification, are permitted provided that the following conditions are met:
+//
+//  1. Redistributions of source code must retain the above copyright notice,
+//     this list of conditions and the following disclaimer.
+//
+//  2. Redistributions in binary form must reproduce the above copyright notice,
+//     this list of conditions and the following disclaimer in the documentation
+//     and/or other materials provided with the distribution.
+//
+//  3. Neither the name of the copyright holder nor the names of its contributors
+//     may be used to endorse or promote products derived from this software
+//     without specific prior written permission.
+//
+//  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+//  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+//  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+//  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+//  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+//  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+//  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+//  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+//  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+//  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+//  POSSIBILITY OF SUCH DAMAGE.
+//
+//  SPDX-License-Identifier: BSD-3-Clause
+//
+//  @@-COPYRIGHT-END-@@
+//
+//==============================================================================
+
+#ifndef AIMET_MAIN_AIMETOPUTILS_H
+#define AIMET_MAIN_AIMETOPUTILS_H
+
+
+#include <DlQuantization/TensorQuantizerOpFacade.h>
+#include <cstdint>
+#include <stdexcept>
+
+
+template <typename T>
+void copyInputTensorsToOutputTensors(const T* inTensor, size_t count, T* outTensor);
+
+
+template <typename T>
+void modeSpecificActionInt(const T* inTensor, size_t count, T* outTensor,
+                           DlQuantization::TensorQuantizerOpFacade* tensorQuantizer,
+                           const DlQuantization::TensorQuantizerOpMode opMode, DlQuantization::TfEncoding* encoding,
+                           const bool useSymmetricEncoding, DlQuantization::IAllocator* allocator)
+{
+    bool useCuda = false;   // Hard code as false until GPU support added
+
+    switch (opMode)
+    {
+    case DlQuantization::TensorQuantizerOpMode::oneShotQuantizeDequantize:
+    {
+        tensorQuantizer->resetEncodingStats();
+        tensorQuantizer->updateStats(inTensor, count, useCuda, allocator);
+        DlQuantization::TfEncoding initial_encoding =
+            tensorQuantizer->computeEncoding(encoding->bw, useSymmetricEncoding);
+        tensorQuantizer->quantizeDequantize(inTensor, count, outTensor, initial_encoding.min, initial_encoding.max,
+                                            encoding->bw, useCuda);
+        // Update encoding object with computed encoding
+        encoding->min    = initial_encoding.min;
+        encoding->max    = initial_encoding.max;
+        encoding->offset = initial_encoding.offset;
+        encoding->delta  = initial_encoding.delta;
+        break;
+    }
+    case DlQuantization::TensorQuantizerOpMode::updateStats:
+    {
+        tensorQuantizer->updateStats(inTensor, count, useCuda, allocator);
+        copyInputTensorsToOutputTensors(inTensor, count, outTensor);
+        break;
+    }
+    case DlQuantization::TensorQuantizerOpMode::quantizeDequantize:
+    {
+        tensorQuantizer->quantizeDequantize(inTensor, count, outTensor, encoding->min, encoding->max, encoding->bw,
+                                            useCuda);
+        break;
+    }
+    case DlQuantization::TensorQuantizerOpMode::passThrough:
+    {
+        copyInputTensorsToOutputTensors(inTensor, count, outTensor);
+        break;
+    }
+    default:
+    {
+        throw std::exception();
+    }
+    }
+}
+
+#endif   // AIMET_MAIN_AIMETOPUTILS_H
diff --git a/TrainingExtensions/onnx/src/QcQuantizeInfo.cpp b/TrainingExtensions/onnx/src/QcQuantizeInfo.cpp
@@ -0,0 +1,57 @@
+//==============================================================================
+//
+//  @@-COPYRIGHT-START-@@
+//
+//  Copyright (c) 2023, Qualcomm Innovation Center, Inc. All rights reserved.
+//
+//  Redistribution and use in source and binary forms, with or without
+//  modification, are permitted provided that the following conditions are met:
+//
+//  1. Redistributions of source code must retain the above copyright notice,
+//     this list of conditions and the following disclaimer.
+//
+//  2. Redistributions in binary form must reproduce the above copyright notice,
+//     this list of conditions and the following disclaimer in the documentation
+//     and/or other materials provided with the distribution.
+//
+//  3. Neither the name of the copyright holder nor the names of its contributors
+//     may be used to endorse or promote products derived from this software
+//     without specific prior written permission.
+//
+//  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+//  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+//  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+//  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+//  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+//  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+//  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+//  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+//  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+//  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+//  POSSIBILITY OF SUCH DAMAGE.
+//
+//  SPDX-License-Identifier: BSD-3-Clause
+//
+//  @@-COPYRIGHT-END-@@
+//
+//==============================================================================
+
+#include "QcQuantizeInfo.h"
+#include <DlQuantization/Quantization.hpp>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(libquant_info, m)
+{
+    pybind11::class_<QcQuantizeInfo>(m, "QcQuantizeInfo")
+        .def(py::init<>())
+        .def_property("tensorQuantizerRef", &QcQuantizeInfo::get_tensor_quantizer,
+                      &QcQuantizeInfo::set_tensor_quantizer)
+        .def_readwrite("encoding", &QcQuantizeInfo::encoding)
+        .def_readwrite("opMode", &QcQuantizeInfo::opMode)
+        .def_readwrite("name", &QcQuantizeInfo::name)
+        .def_readwrite("enabled", &QcQuantizeInfo::enabled)
+        .def_readwrite("useSymmetricEncoding", &QcQuantizeInfo::useSymmetricEncoding);
+}
diff --git a/TrainingExtensions/onnx/src/QcQuantizeInfo.h b/TrainingExtensions/onnx/src/QcQuantizeInfo.h
@@ -0,0 +1,62 @@
+//==============================================================================
+//
+//  @@-COPYRIGHT-START-@@
+//
+//  Copyright (c) 2023, Qualcomm Innovation Center, Inc. All rights reserved.
+//
+//  Redistribution and use in source and binary forms, with or without
+//  modification, are permitted provided that the following conditions are met:
+//
+//  1. Redistributions of source code must retain the above copyright notice,
+//     this list of conditions and the following disclaimer.
+//
+//  2. Redistributions in binary form must reproduce the above copyright notice,
+//     this list of conditions and the following disclaimer in the documentation
+//     and/or other materials provided with the distribution.
+//
+//  3. Neither the name of the copyright holder nor the names of its contributors
+//     may be used to endorse or promote products derived from this software
+//     without specific prior written permission.
+//
+//  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+//  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+//  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+//  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+//  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+//  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+//  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+//  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+//  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+//  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+//  POSSIBILITY OF SUCH DAMAGE.
+//
+//  SPDX-License-Identifier: BSD-3-Clause
+//
+//  @@-COPYRIGHT-END-@@
+//
+//==============================================================================
+
+#pragma once
+
+#include <DlQuantization/IQuantizationEncodingAnalyzer.hpp>
+#include <DlQuantization/QuantizerFactory.hpp>
+#include <string>
+
+
+struct QcQuantizeInfo
+{
+    void set_tensor_quantizer(uint64_t addr)
+    {
+        tensorQuantizerRef = reinterpret_cast<DlQuantization::TensorQuantizerOpFacade*>(addr);
+    }
+    DlQuantization::TensorQuantizer* get_tensor_quantizer() const
+    {
+        return reinterpret_cast<DlQuantization::TensorQuantizer*>(tensorQuantizerRef);
+    }
+    DlQuantization::TensorQuantizerOpFacade* tensorQuantizerRef;
+    DlQuantization::TfEncoding* encoding;
+    DlQuantization::TensorQuantizerOpMode opMode;
+    bool useSymmetricEncoding;
+    bool enabled;
+    std::string name;
+};