Add cuda stream argument for onnx quantize-dequantize

Signed-off-by: Michael Tuttle <[email protected]>
quic · Oct 10, 2023 · 60fbf1f · 60fbf1f
1 parent 626913b
commit 60fbf1f
Show file tree

Hide file tree

Showing 17 changed files with 133 additions and 81 deletions.
diff --git a/ModelOptimizations/DlQuantization/include/DlQuantization/Fp16Quantization.hpp b/ModelOptimizations/DlQuantization/include/DlQuantization/Fp16Quantization.hpp
@@ -49,7 +49,7 @@ namespace DlQuantization
      * @param cnt total size of input tensor
      * @param out pointer to the output tensor
      */
-    void quantizeDequantizeFp16Gpu(const float* in, int cnt, float* out);
+    void quantizeDequantizeFp16Gpu(const float* in, int cnt, float* out, void* stream = nullptr);
 
 }
 

diff --git a/ModelOptimizations/DlQuantization/include/DlQuantization/ITensorQuantizationSim.h b/ModelOptimizations/DlQuantization/include/DlQuantization/ITensorQuantizationSim.h
@@ -52,6 +52,13 @@ class ITensorQuantizationSim
                                           double encodingMin, double encodingMax,
                                           uint8_t bw, RoundingMode roundMode,
                                           bool use_cuda) = 0;
+
+    virtual void quantizeDequantizeTensor(const DTYPE* inputTensorData, size_t inputTensorCount,
+                                          DTYPE* outputTensorData,
+                                          double encodingMin, double encodingMax,
+                                          uint8_t bw, RoundingMode roundMode,
+                                          bool use_cuda, void* stream) = 0;
+
     virtual void quantizeTensor(const DTYPE* inputTensorData, size_t inputTensorCount, DTYPE* outputTensorData,
                                 double encodingMin, double encodingMax, uint8_t bw, RoundingMode roundMode,
                                 bool use_cuda, bool shiftToSigned) = 0;
@@ -112,6 +119,12 @@ class ITensorQuantizationSim
                                                     DTYPE* encodingDelta, DTYPE* encodingOffset,
                                                     RoundingMode roundingMode, bool useCuda) = 0;
 
+    virtual void quantizeDequantizeTensorPerChannel(const DTYPE* inputTensorData, size_t numChannel,
+                                                    size_t numElement, size_t numElementPerChannel,
+                                                    DTYPE* outputTensorData, DTYPE* encodingMin, DTYPE* encodingMax,
+                                                    DTYPE* encodingDelta, DTYPE* encodingOffset,
+                                                    RoundingMode roundingMode, bool useCuda, void* stream) = 0;
+
 };
 
 }   // namespace DlQuantization

diff --git a/ModelOptimizations/DlQuantization/include/DlQuantization/TensorQuantizer.h b/ModelOptimizations/DlQuantization/include/DlQuantization/TensorQuantizer.h
@@ -118,6 +118,9 @@ class TensorQuantizer : public TensorQuantizerOpFacade
     void quantizeDequantize(const float* input, std::size_t tensorSize, float* output,
                             double encodingMin, double encodingMax, unsigned int bitwidth, bool useCuda) override;
 
+    void quantizeDequantize(const float* input, std::size_t tensorSize, float* output, double encodingMin,
+                            double encodingMax, unsigned int bitwidth, bool useCuda, void* stream) override;
+
     /**
      * @brief Convert a tensor from DTYPE to quantized 8-bit packed format
      * @relates quantizeDequantize, except output is stored in 8-bit packed format

diff --git a/ModelOptimizations/DlQuantization/include/DlQuantization/TensorQuantizerOpFacade.h b/ModelOptimizations/DlQuantization/include/DlQuantization/TensorQuantizerOpFacade.h
@@ -90,6 +90,8 @@ class TensorQuantizerOpFacade
     virtual void quantizeDequantize(const float* input, std::size_t tensorSize, float* output,
                                     double encodingMin, double encodingMax, unsigned int bitwidth, bool useCuda) = 0;
 
+    virtual void quantizeDequantize(const float* input, std::size_t tensorSize, float* output, double encodingMin,
+                                    double encodingMax, unsigned int bitwidth, bool useCuda, void* stream) = 0;
     /**
      * Compute the encoding for this tensor using stats collected so far
      */

diff --git a/ModelOptimizations/DlQuantization/src/Fp16Quantization.cpp b/ModelOptimizations/DlQuantization/src/Fp16Quantization.cpp
@@ -43,10 +43,10 @@
 namespace DlQuantization
 {
 
-    void quantizeDequantizeFp16Gpu(const float* in, int cnt, float* out)
+    void quantizeDequantizeFp16Gpu(const float* in, int cnt, float* out, void* stream)
     {
          #ifdef GPU_QUANTIZATION_ENABLED
-                quantizeDequantizeFp16ForGPU(in, cnt, out);
+                quantizeDequantizeFp16ForGPU(in, cnt, out, stream);
          #else
                 throw std::runtime_error("Not compiled for GPU mode.");
          #endif

diff --git a/ModelOptimizations/DlQuantization/src/MainQuantizationClass.cpp b/ModelOptimizations/DlQuantization/src/MainQuantizationClass.cpp
@@ -116,7 +116,7 @@ void MainQuantizationClass<DTYPE>::QuantizeDequantizeActs(const string& layer, L
     for (unsigned int blob_id = 0; blob_id < acts.size(); ++blob_id)
     {
         quantizeDequantize(acts[blob_id], count[blob_id], encoding[blob_id], acts_quantized[blob_id], m_ModeCpuGpu,
-                           ROUND_NEAREST);
+                           ROUND_NEAREST, nullptr);
     }
 }
 
@@ -126,7 +126,7 @@ void MainQuantizationClass<DTYPE>::QuantizeDequantizeParams(int bw, DTYPE* param
                                                             TfEncoding& encoding)
 {
     m_QuantAlgo->NumberDistributionToFxpFormat(bw, params, count, encoding);
-    quantizeDequantize(params, count, encoding, params_quantized, m_ModeCpuGpu, mode_rounding);
+    quantizeDequantize(params, count, encoding, params_quantized, m_ModeCpuGpu, mode_rounding, nullptr);
 }
 
 template <typename DTYPE>

diff --git a/ModelOptimizations/DlQuantization/src/TensorQuantizationSim.cpp b/ModelOptimizations/DlQuantization/src/TensorQuantizationSim.cpp
@@ -95,11 +95,21 @@ void TensorQuantizationSim<DTYPE>::quantizeDequantizeTensor(const DTYPE* inputTe
                                                             DTYPE* outputTensorData, double encodingMin,
                                                             double encodingMax, uint8_t bw, RoundingMode roundingMode,
                                                             bool use_cuda)
+{
+    quantizeDequantizeTensor(inputTensorData, inputTensorCount, outputTensorData, encodingMin, encodingMax,
+                             bw, roundingMode, use_cuda, nullptr);
+}
+
+template <typename DTYPE>
+void TensorQuantizationSim<DTYPE>::quantizeDequantizeTensor(const DTYPE* inputTensorData, size_t inputTensorCount,
+                                                            DTYPE* outputTensorData, double encodingMin,
+                                                            double encodingMax, uint8_t bw, RoundingMode roundingMode,
+                                                            bool use_cuda, void* stream)
 {
     TfEncoding encoding;
     fillEncodingInfo(encoding, bw, encodingMin, encodingMax);
     quantizeDequantize(inputTensorData, inputTensorCount, encoding, outputTensorData, getComputationMode(use_cuda),
-                        roundingMode);
+                        roundingMode, stream);
 }
 
 template <typename DTYPE>
@@ -145,7 +155,7 @@ void TensorQuantizationSim<DTYPE>::quantizeDequantizePerChannelTensor(
     for (uint32_t i = 0; i < splits.size(); ++i) {
         auto& split = splits[i];
         quantizeDequantize(split.data(), split.size(), completeEncodings[i], split.data(), getComputationMode(useCuda),
-                            roundMode);
+                            roundMode, nullptr);
     }
 
     // Concatenate the quantized data back into its original shape.
@@ -269,6 +279,19 @@ void TensorQuantizationSim<DTYPE>::quantizeDequantizeTensorPerChannel(const DTYP
                                                                       DTYPE* encodingMax, DTYPE* encodingDelta,
                                                                       DTYPE* encodingOffset, RoundingMode roundingMode,
                                                                       bool useCuda)
+{
+    quantizeDequantizeTensorPerChannel(inputTensorData, numChannel, numElement, numElementPerChannel, outputTensorData,
+                                       encodingMin, encodingMax, encodingDelta, encodingOffset, roundingMode, useCuda,
+                                       nullptr);
+}
+
+template <typename DTYPE>
+void TensorQuantizationSim<DTYPE>::quantizeDequantizeTensorPerChannel(const DTYPE* inputTensorData, size_t numChannel,
+                                                                      size_t numElement, size_t numElementPerChannel,
+                                                                      DTYPE* outputTensorData, DTYPE* encodingMin,
+                                                                      DTYPE* encodingMax, DTYPE* encodingDelta,
+                                                                      DTYPE* encodingOffset, RoundingMode roundingMode,
+                                                                      bool useCuda, void* stream)
 {
     DlQuantization::ComputationMode cpuGpuMode;
     if (useCuda)
@@ -277,7 +300,8 @@ void TensorQuantizationSim<DTYPE>::quantizeDequantizeTensorPerChannel(const DTYP
         cpuGpuMode = DlQuantization::ComputationMode::COMP_MODE_CPU;
 
     quantizeDequantizePerChannel(inputTensorData, numChannel, numElement, numElementPerChannel, outputTensorData,
-                                 encodingMin, encodingMax, encodingDelta, encodingOffset, cpuGpuMode, roundingMode);
+                                 encodingMin, encodingMax, encodingDelta, encodingOffset, cpuGpuMode, roundingMode,
+                                 stream);
 }
 
 template class TensorQuantizationSim<float>;

diff --git a/ModelOptimizations/DlQuantization/src/TensorQuantizationSim.h b/ModelOptimizations/DlQuantization/src/TensorQuantizationSim.h
@@ -60,6 +60,10 @@ class TensorQuantizationSim : public ITensorQuantizationSim<DTYPE>
                                   double encodingMin, double encodingMax, uint8_t bw, RoundingMode roundMode,
                                   bool use_cuda) override;
 
+    void quantizeDequantizeTensor(const DTYPE* inputTensorData, size_t inputTensorCount, DTYPE* outputTensorData,
+                                  double encodingMin, double encodingMax, uint8_t bw, RoundingMode roundMode,
+                                  bool use_cuda, void* stream) override;
+
     void quantizeTensor(const DTYPE* inputTensorData, size_t inputTensorCount, DTYPE* outputTensorData,
                         double encodingMin, double encodingMax, uint8_t bw, RoundingMode roundMode, bool use_cuda,
                         bool shiftToSigned)
@@ -99,6 +103,11 @@ class TensorQuantizationSim : public ITensorQuantizationSim<DTYPE>
                                             DTYPE* encodingMax, DTYPE* encodingDelta, DTYPE* encodingOffset,
                                             RoundingMode roundingMode, bool useCuda) override;
 
+    void quantizeDequantizeTensorPerChannel(const DTYPE* inputTensorData, size_t numChannel, size_t numElement,
+                                            size_t numElementPerChannel, DTYPE* outputTensorData, DTYPE* encodingMin,
+                                            DTYPE* encodingMax, DTYPE* encodingDelta, DTYPE* encodingOffset,
+                                            RoundingMode roundingMode, bool useCuda, void* stream) override;
+
     inline DlQuantization::ComputationMode getComputationMode(bool use_cuda)
     {
         return (use_cuda ? DlQuantization::ComputationMode::COMP_MODE_GPU

diff --git a/ModelOptimizations/DlQuantization/src/TensorQuantizationSimForPython.cpp b/ModelOptimizations/DlQuantization/src/TensorQuantizationSimForPython.cpp
@@ -75,7 +75,7 @@ py::array_t<float> TensorQuantizationSimForPython::quantizeDequantize(py::array_
 
     _tensorQuantizationSim->quantizeDequantizeTensor(inputDataPtr, inputTensorSize, outputDataPtr,
                                                      encoding.min, encoding.max, bitwidth, roundingMode,
-                                                     use_cuda);
+                                                     use_cuda, nullptr);
 
     return output;
 }

diff --git a/ModelOptimizations/DlQuantization/src/TensorQuantizer.cpp b/ModelOptimizations/DlQuantization/src/TensorQuantizer.cpp
@@ -164,10 +164,16 @@ void TensorQuantizer::computeEncodingFromData(uint8_t bw, const float* data, siz
 
 void TensorQuantizer::quantizeDequantize(const float* input, std::size_t tensorSize, float* output, double encodingMin,
                                          double encodingMax, unsigned int bitwidth, bool useCuda)
+{
+    quantizeDequantize(input, tensorSize, output, encodingMin, encodingMax, bitwidth, useCuda, nullptr);
+}
+
+void TensorQuantizer::quantizeDequantize(const float* input, std::size_t tensorSize, float* output, double encodingMin,
+                                         double encodingMax, unsigned int bitwidth, bool useCuda, void* stream)
 {
     assert(isEncodingValid);
     _tensorQuantizationSim->quantizeDequantizeTensor(input, tensorSize, output, encodingMin, encodingMax, bitwidth,
-                                                     roundingMode, useCuda);
+                                                     roundingMode, useCuda, stream);
 }
 
 void TensorQuantizer::quantizeTensorPacked(const float* input, std::size_t tensorSize, std::vector<uint8_t>& output,

diff --git a/ModelOptimizations/DlQuantization/src/trim_functions.cpp b/ModelOptimizations/DlQuantization/src/trim_functions.cpp
@@ -92,7 +92,7 @@ Lambda parallelize(const uint32_t number_of_threads, Lambda lambda)
 // encoding: TF: rounded
 template <typename DTYPE>
 void quantizeDequantize(const DTYPE* in, int cnt, const TfEncoding& encoding, DTYPE* out,
-                        ComputationMode mode_cpu_gpu, RoundingMode rounding_mode)
+                        ComputationMode mode_cpu_gpu, RoundingMode rounding_mode, void* stream)
 {
     switch (mode_cpu_gpu)
     {
@@ -101,7 +101,7 @@ void quantizeDequantize(const DTYPE* in, int cnt, const TfEncoding& encoding, DT
         break;
     case COMP_MODE_GPU:
 #ifdef GPU_QUANTIZATION_ENABLED
-            quantizeDequantizeGpu(in, cnt, encoding, out, rounding_mode);
+            quantizeDequantizeGpu(in, cnt, encoding, out, rounding_mode, stream);
 #else
         throw runtime_error("Not compiled for GPU mode.");
 #endif
@@ -602,7 +602,7 @@ void dequantizeFromPackedFxpCpu(const uint8_t* input, int cnt,
 template <typename DTYPE>
 void quantizeDequantizePerChannel(const DTYPE* in, int numChannel, int numElement, int numElementPerChannel, DTYPE* out,
                                   DTYPE* encodingMin, DTYPE* encodingMax, DTYPE* encodingDelta, DTYPE* encodingOffset,
-                                  ComputationMode modeCpuGpu, RoundingMode roundingMode)
+                                  ComputationMode modeCpuGpu, RoundingMode roundingMode, void* stream)
 {
     switch (modeCpuGpu)
     {
@@ -613,7 +613,7 @@ void quantizeDequantizePerChannel(const DTYPE* in, int numChannel, int numElemen
     case COMP_MODE_GPU:
 #ifdef GPU_QUANTIZATION_ENABLED
         quantizeDequantizePerChannelGpu(in, numChannel, numElement, numElementPerChannel, out, encodingMin, encodingMax,
-                                        encodingDelta, encodingOffset, roundingMode);
+                                        encodingDelta, encodingOffset, roundingMode, stream);
 #else
         throw runtime_error("Not compiled for GPU mode.");
 #endif
@@ -643,10 +643,10 @@ void quantizeDequantizePerChannelCpu(const DTYPE* in, int numChannel, int numEle
 
 // Explicit instantiations
 template void quantizeDequantize(const double* in, int cnt, const TfEncoding& encoding, double* out,
-                                 ComputationMode mode_cpu_gpu, RoundingMode rounding_mode);
+                                 ComputationMode mode_cpu_gpu, RoundingMode rounding_mode, void* stream);
 
 template void quantizeDequantize(const float* in, int cnt, const TfEncoding& encoding, float* out,
-                                 ComputationMode mode_cpu_gpu, RoundingMode rounding_mode);
+                                 ComputationMode mode_cpu_gpu, RoundingMode rounding_mode, void* stream);
 
 template void quantizeToFxp(const double* in, int cnt, const TfEncoding& encoding, double* out,
                             ComputationMode mode_cpu_gpu, RoundingMode rounding_mode, bool shiftToSigned);
@@ -669,9 +669,9 @@ template void dequantizeFromPackedFxp(const uint8_t* input, int cnt,
 
 template void quantizeDequantizePerChannel(const float* in, int numChannel, int numElement, int numElementPerChannel, float* out,
                                            float* encodingMin, float* encodingMax, float* encodingDelta, float* encodingOffset,
-                                           ComputationMode modeCpuGpu, RoundingMode roundingMode);
+                                           ComputationMode modeCpuGpu, RoundingMode roundingMode, void* stream);
 template void quantizeDequantizePerChannel(const double* in, int numChannel, int numElement, int numElementPerChannel, double* out,
                                            double* encodingMin, double* encodingMax, double* encodingDelta, double* encodingOffset,
-                                           ComputationMode modeCpuGpu, RoundingMode roundingMode);
+                                           ComputationMode modeCpuGpu, RoundingMode roundingMode, void* stream);
 
 }   // End of namespace DlQuantization