Merge pull request #189 from Xilinx/matthias.backport_avgpool

Backport: [ONNX] Fixes Issue with Dynamic Dims in GlobalAveragePool -> Torch Conversion
Xilinx · Jul 4, 2024 · e6c7403 · e6c7403
2 parents 3cbdae8 + 7cdea15
commit e6c7403
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -791,9 +791,17 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
         Value cstOne = rewriter.create<Torch::ConstantIntOp>(
             binder.getLoc(), rewriter.getI64IntegerAttr(1));
         for (unsigned i = 2; i < inputRank; i++) {
-          int64_t kernelSize = inputShape[i] - resultShape[i] + 1;
-          cstKernel.push_back(rewriter.create<Torch::ConstantIntOp>(
-              binder.getLoc(), rewriter.getI64IntegerAttr(kernelSize)));
+          if (inputShape[i] == Torch::kUnknownSize) {
+            Value dim = rewriter.create<Torch::ConstantIntOp>(
+                binder.getLoc(), rewriter.getI64IntegerAttr(i));
+            Value inputDimSize = rewriter.create<Torch::AtenSizeIntOp>(
+                binder.getLoc(), operand, dim);
+            cstKernel.push_back(inputDimSize);
+          } else {
+            int64_t kernelSize = inputShape[i] - resultShape[i] + 1;
+            cstKernel.push_back(rewriter.create<Torch::ConstantIntOp>(
+                binder.getLoc(), rewriter.getI64IntegerAttr(kernelSize)));
+          }
           cstPadding.push_back(cstZero);
           cstStrides.push_back(cstOne);
         }

diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -1661,8 +1661,6 @@
     "PermuteNegativeIndexModule_basic",
 
     # Failure - incorrect numerics
-    "AdaptiveAvgPool1dUnitOutputSizeDynamicModule_basic",
-    "AdaptiveAvgPool2dUnitOutputSizeDynamicModule_basic",
     "ElementwiseAtan2TensorIntModule_basic",
     "ElementwiseLog10IntModule_basic",
     "ElementwiseLog2IntModule_basic",
@@ -1672,7 +1670,6 @@
     "HardsigmoidModule_basic",
     "HardsigmoidRandomModule_basic",
     "PixelShuffleModuleStaticRank4Float32_basic",
-    "ResNet18Module_basic",
     "SliceCopyEndGreaterThanDimSize_Module_basic",
     "SliceCopyNegative_Module_basic",
     "SliceCopyNonZeroDim_Module_basic",