add uncovered cases in view lowering (#2524)

removes unecessary checks from empty strided

Author: dan-garvey

e2e_testing/xfail_sets.py

@@ -154,6 +154,9 @@
     'BoolFloatConstantModule_basic',
     'BoolIntConstantModule_basic',
 
+    # ERROR: torch._dynamo.exc.Unsupported: torch.* op returned non-Tensor int call_function aten.size
+    "ViewSizeFromOtherTensor_basic",
+
     # ERROR: torch._dynamo.exc.Unsupported: torch.* op returned non-Tensor bool call_function aten.__contains__
     'ContainsIntList_False',
     'ContainsIntList_True',

lib/Conversion/TorchToLinalg/DataMovement.cpp

@@ -377,6 +377,11 @@ class ConvertAtenViewOp : public OpConversionPattern<AtenViewOp> {
     // collapsed. Note this may technically not always be true.
     // TODO: think of a way better way to at least detect when this assumption
     // is violated for the cases of dynamic dimensions.
+    bool inputHasOneDynDim = llvm::count(inputShape, kUnknownSize) == 1;
+    bool outputHasOneDynDim = llvm::count(outputShape, kUnknownSize) == 1;
+    bool singleDynDimsAreEqual =
+    inputHasOneDynDim && outputHasOneDynDim &&
+    productReduce(inputShape) == productReduce(outputShape);
     SmallVector<std::pair<int64_t, int64_t>> unchangedDims;
     for (auto [outputDim, outputDimSize] :
          llvm::enumerate(outputSizeTorchInt)) {
@@ -385,6 +390,14 @@ class ConvertAtenViewOp : public OpConversionPattern<AtenViewOp> {
       if (matchPattern(outputDimSize,
                        m_TorchTensorSizeInt(op.getSelf(), &inputDim))) {
         unchangedDims.push_back(std::make_pair(inputDim, outputDim));
+      } else if (singleDynDimsAreEqual &&
+                 outputShape[outputDim] == kUnknownSize) {
+        // If the input and output have a single dynamic dimension and the
+        // product of the other dimensions is the same, then we know that the
+        // dynamic dimension is unchanged.
+        inputDim = std::distance(inputShape.begin(),
+                                 llvm::find(inputShape, kUnknownSize));
+        unchangedDims.push_back(std::make_pair(inputDim, outputDim));
       }
     }
     // Mark the end of the input/output shapes

python/torch_mlir_e2e_test/test_suite/reshape_like.py

@@ -430,6 +430,26 @@ def ViewFlattenAndExpandModule_basic(module, tu: TestUtils):
 
 # ==============================================================================
 
+class ViewSizeFromOtherTensor(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([1, -1], torch.float32, True),
+        ([1, -1, 10], torch.float32, True),
+    ])
+
+    def forward(self, x, y):
+        return torch.ops.aten.view(y, (torch.ops.aten.size(x, 1), 10))
+
+@register_test_case(module_factory=lambda: ViewSizeFromOtherTensor())
+def ViewSizeFromOtherTensor_basic(module, tu: TestUtils):
+    module.forward(tu.rand(1, 7), tu.rand(1, 7, 10))
+
+# ==============================================================================
+
 class UnsafeViewExpandModule(torch.nn.Module):
     def __init__(self):
         super().__init__()

test/Conversion/TorchToLinalg/view.mlir

@@ -73,4 +73,4 @@ func.func @torch.aten.view$expandInferredDim(%arg0: !torch.vtensor<[2,6],f32>) -
     %0 = torch.prim.ListConstruct %int3, %int2, %int-1 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
     %1 = torch.aten.view %arg0, %0 : !torch.vtensor<[2,6],f32>, !torch.list<int> -> !torch.vtensor<[3,2,2],f32>
     return %1 : !torch.vtensor<[3,2,2],f32>
-  }
+  }