feat(//core/conversion/conversionctx): Make op precision available at

conversion time through ctx

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/conversion/conversion.cpp

@@ -160,6 +160,10 @@ void AddInputs(ConversionCtx* ctx,
     TRTORCH_CHECK(profile->isValid(), "Optimization profile is invalid, please check the input range provided (conversion.AddInputs)");
 
     ctx->cfg->addOptimizationProfile(profile);
+    // TODO: Enable in TRT 7.1
+    // if (ctx->op_precision == nvinfer1::DataType::kINT8) {
+    //     ctx->cfg->setCalibrationProfile(profile);
+    // }
 }
 
 void MarkOutputs(ConversionCtx* ctx, at::ArrayRef<const torch::jit::Value*> outputs) {

core/conversion/conversionctx/ConversionCtx.cpp

@@ -60,6 +60,7 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
         input_type = nvinfer1::DataType::kFLOAT;
         break;
     }
+    op_precision = settings.op_precision;
 
     if (settings.refit) {
         cfg->setFlag(nvinfer1::BuilderFlag::kREFIT);

core/conversion/conversionctx/ConversionCtx.h

@@ -47,6 +47,7 @@ struct ConversionCtx {
     nvinfer1::INetworkDefinition* net;
     nvinfer1::IBuilderConfig* cfg;
     nvinfer1::DataType input_type;
+    nvinfer1::DataType op_precision;
     BuilderSettings settings;
     util::logging::TRTorchLogger logger;
     // Pointers to data that needs to remain alive until conversion is done