[onnx] handle dynamic padSize tensor onnx.Pad

llvm · Apr 23, 2024 · 897d0f0 · 897d0f0
1 parent c1967b6
commit 897d0f0
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -1319,9 +1319,18 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
                                              "expect 1-d pad tensor");
 
         int64_t padsSize = padsShape[0];
-        if (padsSize == Torch::kUnknownSize)
-          return rewriter.notifyMatchFailure(binder.op,
-                                             "pad length is unknown");
+        if (padsSize == Torch::kUnknownSize) {
+          // As per onnx.Pad documentation, padSize = 2*num_data_axes
+          // (if axes param not passed). Need to be updated when adding
+          // support for `axes` param.
+          auto dataOpTy = data.getType().cast<Torch::ValueTensorType>();
+          TensorType dataTensor = dataOpTy.toBuiltinTensor();
+          if (!dataTensor || !dataTensor.hasRank())
+            return rewriter.notifyMatchFailure(
+                binder.op, "pad length unknown and data operand unranked");
+          int64_t dataRank = dataTensor.getRank();
+          padsSize = 2 * dataRank;
+        }
 
         Value constantValue;
         if (binder.getNumOperands() >= 3) {