[MLIR][TORCH] Support unknown sized input during decomposition of aten.as_strided op (#4303)

#4269  adds the e2e support for the aten.as_strided op by
decomposing it into a series of other torch operations. This change
extents #4269 to relax a restrictive check and support unknown sized
inputs during decomposition

Author: vinitdeodhar

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -12664,8 +12664,8 @@ class DecomposeAtenAsStridedOp : public OpRewritePattern<AtenAsStridedOp> {
     Value input = op.getSelf();
     auto inputType = dyn_cast<BaseTensorType>(input.getType());
 
-    if (!inputType || !inputType.hasSizes() || !inputType.areAllSizesKnown())
-      return rewriter.notifyMatchFailure(op, "input must have known sizes");
+    if (!inputType || !inputType.hasSizes())
+      return rewriter.notifyMatchFailure(op, "input must have sizes");
 
     SmallVector<int64_t> sizesInts;
     if (!matchPattern(op.getSize(), m_TorchListOfConstantInts(sizesInts)))
@@ -12695,8 +12695,13 @@ class DecomposeAtenAsStridedOp : public OpRewritePattern<AtenAsStridedOp> {
       // If the input is not a 1-d tensor, we need to flatten it
       // to a 1D tensor before applying the strided indexing.
       int64_t flattenedInputSize = 1;
-      for (int64_t size : inputSizes)
+      for (int64_t size : inputSizes) {
+        if (size == kUnknownSize) {
+          flattenedInputSize = kUnknownSize;
+          break;
+        }
         flattenedInputSize *= size;
+      }
 
       auto flattenedInputTy =
           cast<BaseTensorType>(inputType.getWithSizesAndDtype(

projects/pt1/e2e_testing/xfail_sets.py

@@ -984,6 +984,7 @@
     "MaxPool2dCeilModeFullDimIndivisibleByStrideModule_basic",
     "AtenAsStridedModule_basic",
     "AtenAsStridedNoStorageOffsetModule_basic",
+    "AtenAsStridedUnknownSizeModule_basic",
     # error: argument must be a memref of f32, f64, i32, i64, i8, i1, c32, c64, but got 'memref<3x5xbf16>'
     "ElementwiseClampMaxModule_bfloat16",
     "ElementwiseClampMinModule_bfloat16",
@@ -3991,6 +3992,7 @@
     "ReplicationPad3dModuleSingleIntPad_basic",
     "AtenAsStridedModule_basic",
     "AtenAsStridedNoStorageOffsetModule_basic",
+    "AtenAsStridedUnknownSizeModule_basic",
     "ChunkListUnpackDynamic_Module_basic",
     "ChunkListUnpackUnevenDynamic_Module_basic",
     "ChunkListUnpackUneven_Module_basic",

projects/pt1/python/torch_mlir_e2e_test/test_suite/basic.py

@@ -6974,3 +6974,23 @@ def forward(self, x):
 @register_test_case(module_factory=lambda: AtenAsStridedNoStorageOffsetModule())
 def AtenAsStridedNoStorageOffsetModule_basic(module, tu: TestUtils):
     module.forward(torch.randn(12, 13))
+
+
+class AtenAsStridedUnknownSizeModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, 13], torch.float32, True),
+        ]
+    )
+    def forward(self, x):
+        return torch.ops.aten.as_strided(x, size=(3, 4), stride=(2, 2))
+
+
+@register_test_case(module_factory=lambda: AtenAsStridedUnknownSizeModule())
+def AtenAsStridedUnknownSizeModule_basic(module, tu: TestUtils):
+    module.forward(torch.randn(12, 13))