Rewrote solution via decomposition

-Co-author: @ivanamitreski

Author: amemov

lib/Conversion/TorchToLinalg/Linear.cpp

@@ -1673,80 +1673,6 @@ struct ConvertAtenFftRfftOp final : OpConversionPattern<AtenFftRfftOp> {
 
 } // namespace
 
-namespace {
-class ConvertAtenOuterOp : public OpConversionPattern<AtenOuterOp> {
-public:
-  using OpConversionPattern::OpConversionPattern;
-  LogicalResult
-  matchAndRewrite(AtenOuterOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-
-    Location loc = op->getLoc();
-    Value lhs = adaptor.getSelf();
-    Value rhs = adaptor.getVec2();
-
-    if (failed(verifyLinalgCompatibleTypes(op, rewriter))) {
-      return failure();
-    }
-    auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
-    auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
-
-    if (!lhsType || !rhsType)
-      return rewriter.notifyMatchFailure(op,
-                                         "outer: expected ranked tensor types");
-    if (lhsType.getRank() != 1 || rhsType.getRank() != 1)
-      return rewriter.notifyMatchFailure(
-          op, "outer: expected 1D tensors for outer op lowering");
-
-    Value lhsDim = getDimOp(rewriter, loc, lhs, 0);
-    Value rhsDim = getDimOp(rewriter, loc, rhs, 0);
-    Type elementType = lhsType.getElementType();
-    Type newResultType = getTypeConverter()->convertType(op.getType());
-
-    // Create a zero-initialized tensor with shape [lhsDim, rhsDim]
-    SmallVector<OpFoldResult> resultShape =
-        getAsOpFoldResult(ValueRange{lhsDim, rhsDim});
-    Value initTensor =
-        rewriter.create<tensor::EmptyOp>(loc, resultShape, elementType);
-
-    // Set up affine indexing maps:
-    // We create a 2D loop iteration space. For the lhs, we use the first index
-    // (i), for the rhs, the second index (j), and for the result, both (i, j).
-    AffineMap mapLhs =
-        AffineMap::get(2, /*symbolCount=*/0, {rewriter.getAffineDimExpr(0)},
-                       rewriter.getContext());
-    AffineMap mapRhs =
-        AffineMap::get(2, /*symbolCount=*/0, {rewriter.getAffineDimExpr(1)},
-                       rewriter.getContext());
-    AffineMap mapOut =
-        AffineMap::getMultiDimIdentityMap(2, rewriter.getContext());
-
-    SmallVector<utils::IteratorType, 2> iteratorTypes = {
-        utils::IteratorType::parallel, utils::IteratorType::parallel};
-
-    Value outerProd =
-        rewriter
-            .create<linalg::GenericOp>(
-                loc, initTensor.getType(),
-                /*inputs=*/ValueRange{lhsDim, rhsDim},
-                /*outputs=*/initTensor,
-                /*indexingMaps=*/
-                SmallVector<AffineMap, 3>{mapLhs, mapRhs, mapOut},
-                /*iteratortType=*/iteratorTypes,
-                [&](OpBuilder &b, Location loc, ValueRange args) {
-                  Value lhsElem = args[0];
-                  Value rhsElem = args[1];
-                  Value mult = b.create<arith::MulFOp>(loc, lhsElem, rhsElem);
-                  b.create<linalg::YieldOp>(loc, mult);
-                })
-            .getResult(0);
-
-    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, newResultType, outerProd);
-    return success();
-  }
-};
-} // namespace
-
 void mlir::torch::torch_to_linalg::populateLinearPatternsAndLegality(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     ConversionTarget &target) {
@@ -1763,6 +1689,4 @@ void mlir::torch::torch_to_linalg::populateLinearPatternsAndLegality(
   patterns.add<ConvertAtenConvolutionOp>(typeConverter, context);
   target.addIllegalOp<AtenFftRfftOp>();
   patterns.add<ConvertAtenFftRfftOp>(typeConverter, context);
-  target.addIllegalOp<AtenOuterOp>();
-  patterns.add<ConvertAtenOuterOp>(typeConverter, context);
 }

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -1894,6 +1894,59 @@ class DecomposeAtenAtleast1dOp : public OpRewritePattern<AtenAtleast1dOp> {
 };
 } // namespace
 
+// Decompose 'aten.outer' into 'aten.unsqueeze', 'aten.matmul'
+
+namespace {
+class DecomposeAtenOuterOp : public OpRewritePattern<AtenOuterOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenOuterOp op,
+                                PatternRewriter &rewriter) const override {
+
+    Location loc = op.getLoc();
+    Value input = op.getSelf();
+    Value vec2 = op.getVec2();
+    Type opType = op.getType();
+
+    auto inputType = cast<BaseTensorType>(input.getType());
+    auto vec2Type = cast<BaseTensorType>(vec2.getType());
+
+    // Check if both tensors are 1-dimensional
+    SmallVector<int64_t> inputShape(inputType.getSizes());
+    SmallVector<int64_t> vec2Shape(vec2Type.getSizes());
+
+    if (inputShape.size() == 1 && vec2Shape.size() == 1) {
+
+      Value one = rewriter.create<Torch::ConstantIntOp>(
+          loc, rewriter.getI64IntegerAttr(1)); // Dimension index
+      SmallVector<int64_t, 2> inputMatrixShape = {inputShape[0], 1};
+      Type inputMatrixType = inputType.getWithSizesAndDtype(
+          inputMatrixShape, inputType.getOptionalDtype());
+
+      Value inputMatrix =
+          rewriter.create<AtenUnsqueezeOp>(loc, inputMatrixType, input, one);
+
+      Value zero = rewriter.create<Torch::ConstantIntOp>(
+          loc, rewriter.getI64IntegerAttr(0));
+      SmallVector<int64_t, 2> vec2MatrixShape = {1, vec2Shape[0]};
+      Type vec2MatrixType = vec2Type.getWithSizesAndDtype(
+          vec2MatrixShape, vec2Type.getOptionalDtype());
+
+      Value vec2Matrix =
+          rewriter.create<AtenUnsqueezeOp>(loc, vec2MatrixType, vec2, zero);
+
+      rewriter.replaceOpWithNewOp<AtenMatmulOp>(op, opType, inputMatrix,
+                                                vec2Matrix);
+      return success();
+    } else {
+      return failure();
+    }
+
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 // Decompose aten.atleast_2d into: aten.reshape. See
 // https://github.com/pytorch/pytorch/blob/9a8ab778d34bd24c5caceb340837483decc4c311/torch/_refs/__init__.py#L2604
@@ -11591,6 +11644,7 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAtenEmptyLikeOp>(patterns);
     addPatternIfTargetOpIsIllegal<
         DecomposeConstantTensorAllocLikeOp<AtenOnesLikeOp, 1>>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenOuterOp>(patterns);
     addPatternIfTargetOpIsIllegal<
         DecomposeConstantTensorAllocLikeOp<AtenZerosLikeOp, 0>>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenStackOp>(patterns);

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -379,6 +379,7 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
   target.addIllegalOp<AtenSoftshrinkOp>();
   target.addIllegalOp<AtenEmptyLikeOp>();
   target.addIllegalOp<AtenOnesLikeOp>();
+  target.addIllegalOp<AtenOuterOp>();
   target.addIllegalOp<AtenZerosLikeOp>();
   target.addIllegalOp<AtenStackOp>();
   target.addIllegalOp<AtenHstackOp>();

projects/pt1/e2e_testing/xfail_sets.py

@@ -3810,6 +3810,7 @@
 }
 
 ONNX_TOSA_XFAIL_SET = {
+    "AtenOuter_basic",
     "AtenFftRfft2DLastDim_basic",
     "AtenFftRfft2DMiddleDim_basic",
     "AtenNonzero1DDynamicModule_basic",