[MLIR][TORCH] Add E2E support for aten.squeeze op

This commit adds lowering of `aten.Squeeze` op into
`linalg.TensorCollapseShape` op. The size 1 dynamic dimensions are not
handled as a part of this commit.

Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>

Author: Shukla-Gaurav

e2e_testing/torchscript/main.py

@@ -41,6 +41,7 @@
 from . import matmul
 from . import view
 from . import scalar
+from . import squeeze
 
 def _get_argparse():
     config_choices = ['native_torch', 'torchscript', 'refbackend', 'tosa', 'external']

e2e_testing/torchscript/squeeze.py

@@ -0,0 +1,121 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# Also available under a BSD-style license. See LICENSE.
+
+import torch
+
+from torch_mlir_e2e_test.torchscript.framework import TestUtils
+from torch_mlir_e2e_test.torchscript.registry import register_test_case
+from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+
+# ==============================================================================
+
+
+class SqueezeStaticModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([1, 7, 1, 3, 1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.squeeze(a)
+
+
+@register_test_case(
+    module_factory=lambda: SqueezeStaticModule())
+def SqueezeModule_static(module, tu: TestUtils):
+    module.forward(tu.rand(1, 7, 1, 3, 1))
+
+
+# ==============================================================================
+
+
+class SqueezeDynamicModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([1, -1, 1, 384, -1, 1, 1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.squeeze(a)
+
+
+@register_test_case(
+    module_factory=lambda: SqueezeDynamicModule())
+def SqueezeModule_dynamic(module, tu: TestUtils):
+    module.forward(tu.rand(1, 8, 1, 384, 12, 1, 1))
+
+
+# ==============================================================================
+
+
+class SqueezeNoUnitDimModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([4, -1, -1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.squeeze(a)
+
+
+@register_test_case(
+    module_factory=lambda: SqueezeNoUnitDimModule())
+def SqueezeModule_noUnitDim(module, tu: TestUtils):
+    module.forward(tu.rand(4, 2, 3))
+
+
+# ==============================================================================
+
+
+class SqueezeAllUnitDimModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([1, 1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.squeeze(a)
+
+
+@register_test_case(
+    module_factory=lambda: SqueezeAllUnitDimModule())
+def SqueezeModule_allUnitDim(module, tu: TestUtils):
+    module.forward(tu.rand(1, 1))
+
+
+# ==============================================================================
+
+
+class SqueezeBroadcastModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+        ([], torch.float32, True),
+    ])
+    def forward(self, a, b):
+        return a * b.squeeze()
+
+
+@register_test_case(
+    module_factory=lambda: SqueezeBroadcastModule())
+def SqueezeModule_broadcast(module, tu: TestUtils):
+    module.forward(tu.rand(4, 3), tu.rand())
+

include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td

@@ -1452,6 +1452,20 @@ def Torch_AtenUnsqueezeOp : Torch_Op<"aten.unsqueeze", [
   let assemblyFormat = "$self `,` $dim attr-dict `:` type($self) `,` type($dim) `->` type($result)";
 }
 
+def Torch_AtenSqueezeOp : Torch_Op<"aten.squeeze", [
+    AllowsTypeRefinement
+  ]> {
+  let summary = "Generated op for `aten::squeeze : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
+  let hasFolder = 1;
+}
+
 def Torch_AtenFlattenUsingIntsOp : Torch_Op<"aten.flatten.using_ints", [
     AllowsTypeRefinement
   ]> {

lib/Conversion/TorchToLinalg/TorchToLinalg.cpp

@@ -2388,6 +2388,97 @@ class ConvertAtenViewOp : public OpConversionPattern<AtenViewOp> {
 };
 } // namespace
 
+namespace {
+class ConvertAtenSqueezeOp : public OpConversionPattern<AtenSqueezeOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(AtenSqueezeOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
+      return failure();
+    Location loc = op.getLoc();
+    Value input = adaptor.self();
+    auto inputType = input.getType().cast<RankedTensorType>();
+    int64_t inputRank = inputType.getRank();
+    TypeConverter *typeConverter = getTypeConverter();
+    auto resultType =
+        typeConverter->convertType(op.getType()).cast<RankedTensorType>();
+    int64_t resultRank = resultType.getRank();
+
+    if (inputRank == 0) {
+      return rewriter.notifyMatchFailure(
+          op, "zero input rank should have been handled by the folder");
+    }
+
+    // In case the operand tensor type is statically shaped with all dimensions
+    // being unit extent, it will be collapsed to a 0-D tensor.
+    if (resultRank == 0) {
+      SmallVector<ReassociationIndices> reassociation;
+      rewriter.replaceOpWithNewOp<linalg::TensorCollapseShapeOp>(
+          op, resultType, input, reassociation);
+      return success();
+    }
+
+    // All the static size-1 dimensions at the beginning(going from higher to
+    // lower dimensions) will be collapsed into the first dynamic or first non
+    // size-1 static dimension. All the other static size-1 dimensions will be
+    // collapsed into its previous dynamic or non size-1 static dimension.
+    SmallVector<ReassociationIndices> reassociation(resultRank);
+    bool isSqueezed = false;
+    int64_t headOnesCount = 0;
+    while (headOnesCount < inputRank &&
+           inputType.getDimSize(headOnesCount) == 1) {
+      isSqueezed = true;
+      reassociation[0].push_back(headOnesCount++);
+    }
+
+    // TODO: Add support for size-1 dynamic dimensions.
+    Value one = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getIntegerAttr(rewriter.getIndexType(), 1));
+    int64_t j = -1;
+    for (auto i : llvm::seq<int64_t>(headOnesCount, inputRank)) {
+      if (inputType.isDynamicDim(i)) {
+        // Make sure that size-1 dynamic dimension does not exist.
+        Value dimSize = getDimOp(rewriter, loc, input, i);
+        Value dimSizeNotOne = rewriter.create<arith::CmpIOp>(
+            loc, arith::CmpIPredicate::ne, dimSize, one);
+        rewriter.create<AssertOp>(
+            loc, dimSizeNotOne,
+            rewriter.getStringAttr(
+                "unimplemented: size 1 dynamic dimension is not supported"));
+        ++j;
+      } else if (inputType.getDimSize(i) != 1) {
+        ++j;
+      } else {
+        // `isSqueezed` checks if the operand tensor type contains at least one
+        // unit dimension.
+        isSqueezed = true;
+      }
+      if (j == resultRank)
+        break;
+      reassociation[j].push_back(i);
+    }
+
+    // Make sure that result type rank is compatible with the squeezed size.
+    if (j != resultRank - 1)
+      return rewriter.notifyMatchFailure(
+          op, "expected output size mismatches with the result type rank");
+
+    if (isSqueezed) {
+      rewriter.replaceOpWithNewOp<linalg::TensorCollapseShapeOp>(
+          op, resultType, input, reassociation);
+
+    } else {
+      // If the operand tensor type does not have any unit dimension,
+      // `aten.squeeze` will behave as an identity operation.
+      rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, input);
+    }
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class ConvertAtenUnsqueezeOp : public OpConversionPattern<AtenUnsqueezeOp> {
 public:
@@ -3057,6 +3148,8 @@ class ConvertTorchToLinalg
         AtenRsubScalarOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
         AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp>();
     patterns.add<ConvertElementwiseOp>(typeConverter, context);
+    target.addIllegalOp<AtenSqueezeOp>();
+    patterns.add<ConvertAtenSqueezeOp>(typeConverter, context);
     target.addIllegalOp<AtenUnsqueezeOp>();
     patterns.add<ConvertAtenUnsqueezeOp>(typeConverter, context);
     target.addIllegalOp<AtenConv2dOp>();

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -450,6 +450,18 @@ OpFoldResult AtenNeBoolOp::fold(ArrayRef<Attribute> operands) {
   return IntegerAttr::get(IntegerType::get(getContext(), 1), a != b);
 }
 
+//===----------------------------------------------------------------------===//
+// AtenSqueezeOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AtenSqueezeOp::fold(ArrayRef<Attribute> operands) {
+  if (auto tensorType = getOperand().getType().dyn_cast<BaseTensorType>()) {
+    if (tensorType.hasSizes() && tensorType.getSizes().size() == 0)
+      return getOperand();
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // AtenDimOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Torch/Transforms/MaximizeValueSemantics.cpp

@@ -89,7 +89,7 @@ class RewriteViewLikeSubgraph
       Operation *op = workList.pop_back_val();
       if (auto copyToValueTensor = dyn_cast<CopyToValueTensorOp>(op)) {
         copyToValueTensorOps.push_back(copyToValueTensor);
-      } else if (isa<AtenUnsqueezeOp, AtenFlattenUsingIntsOp,
+      } else if (isa<AtenSqueezeOp, AtenUnsqueezeOp, AtenFlattenUsingIntsOp,
                      AtenTransposeIntOp, TensorStaticInfoCastOp,
                      AtenBroadcastToOp, AtenToDtypeOp, AtenContiguousOp,
                      AtenPermuteOp, AtenViewOp, AtenExpandOp,

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -300,6 +300,8 @@ class TypeAnalyzer : public ForwardDataFlowAnalysis<ValueKnowledge> {
       return visitAtenLerpTensorOp(lerpTensor, operands);
     } else if (auto flatten = dyn_cast<AtenFlattenUsingIntsOp>(op)) {
       return visitAtenFlattenUsingIntsOp(flatten, operands);
+    } else if (auto squeeze = dyn_cast<AtenSqueezeOp>(op)) {
+      return visitAtenSqueezeOp(squeeze, operands);
     } else if (auto unsqueeze = dyn_cast<AtenUnsqueezeOp>(op)) {
       return visitAtenUnsqueezeOp(unsqueeze, operands);
     } else if (auto arange = dyn_cast<AtenArangeOp>(op)) {
@@ -466,6 +468,9 @@ class TypeAnalyzer : public ForwardDataFlowAnalysis<ValueKnowledge> {
       AtenFlattenUsingIntsOp op,
       ArrayRef<LatticeElement<ValueKnowledge> *> operands);
   ChangeResult
+  visitAtenSqueezeOp(AtenSqueezeOp op,
+                     ArrayRef<LatticeElement<ValueKnowledge> *> operands);
+  ChangeResult
   visitAtenUnsqueezeOp(AtenUnsqueezeOp op,
                        ArrayRef<LatticeElement<ValueKnowledge> *> operands);
 
@@ -880,6 +885,25 @@ ChangeResult TypeAnalyzer::visitAtenFlattenUsingIntsOp(
   return getLatticeElement(op.getResult()).join(knowledge);
 }
 
+ChangeResult TypeAnalyzer::visitAtenSqueezeOp(
+    AtenSqueezeOp op, ArrayRef<LatticeElement<ValueKnowledge> *> operands) {
+  auto operand = operands[0]->getValue();
+  auto knowledge =
+      ValueKnowledge::getNotNonePessimisticValueState(op.getContext());
+  knowledge.dtype = operand.dtype;
+  if (operand.hasSizes) {
+    int64_t inputRank = operand.sizes.size();
+    knowledge.hasSizes = true;
+    // `knowledge.sizes` will be empty when either `inputRank` is 0 or operand
+    // tensor type is statically shaped with all dimensions being unit.
+    // Note: size-1 dynamic dimensions are not supported yet.
+    for (auto i = 0; i < inputRank; i++)
+      if (operand.sizes[i] != 1)
+        knowledge.sizes.push_back(operand.sizes[i]);
+  }
+  return getLatticeElement(op.getResult()).join(knowledge);
+}
+
 ChangeResult TypeAnalyzer::visitAtenUnsqueezeOp(
     AtenUnsqueezeOp op, ArrayRef<LatticeElement<ValueKnowledge> *> operands) {
   auto operand = operands[0]->getValue();

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -521,6 +521,7 @@ def emit_with_mutating_variants(key, **kwargs):
 
         # Misc tensor ops.
         emit("aten::unsqueeze : (Tensor, int) -> (Tensor)")
+        emit("aten::squeeze : (Tensor) -> (Tensor)", has_folder=True)
         emit("aten::flatten.using_ints : (Tensor, int, int) -> (Tensor)")
         emit("aten::dim : (Tensor) -> (int)", has_folder=True)
         emit("aten::size : (Tensor) -> (int[])", has_canonicalizer=True)

test/Dialect/Torch/canonicalize.mlir

@@ -594,3 +594,11 @@ func @torch.prim.TupleIndex$out_of_bound(%t0: !torch.tensor, %t1: !torch.tensor,
     %1 = torch.prim.TupleIndex %0, %int3 : !torch.tuple<!torch.tensor, !torch.tensor, !torch.tensor>, !torch.int -> !torch.tensor
     return %1 : !torch.tensor
 }
+
+// CHECK-LABEL:   func @torch.aten.squeeze$zero_rank(
+// CHECK-SAME:            %[[ARG:.*]]: !torch.tensor<[],f32>) -> !torch.tensor<[],f32> {
+// CHECK-NEXT:      return %[[ARG]] : !torch.tensor<[],f32>
+func @torch.aten.squeeze$zero_rank(%arg0: !torch.tensor<[],f32>) -> !torch.tensor<[],f32> {
+  %0 = torch.aten.squeeze %arg0 : !torch.tensor<[],f32> -> !torch.tensor<[],f32>
+  return %0 : !torch.tensor<[],f32>
+}