Implementation of torch-to-linalg lowering of AtenOuterOp

lib/Conversion/TorchToLinalg/Linear.cpp

@@ -1673,6 +1673,80 @@ 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) {
@@ -1689,4 +1763,6 @@ 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);
 }

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

@@ -918,3 +918,61 @@ def forward(self, a, b):
 @register_test_case(module_factory=lambda: AtenLinalgCrossDynamic())
 def AtenLinalgCrossDynamic_basic(module, tu: TestUtils):
     module.forward(tu.rand(4, 3, 1, 6), tu.rand(4, 3, 7, 1))
+
+
+# ==============================================================================
+
+
+class AtenOuter(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([3], torch.float32, True),
+            ([3], torch.float32, True),
+        ]
+    )
+    def forward(self, lhs, rhs):
+        return torch.outer(lhs, rhs)
+
+
+@register_test_case(module_factory=lambda: AtenOuter())
+def AtenOuter_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3), tu.rand(3))
+
+
+# ==============================================================================
+
+
+class AtenOuterDynamic(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1], torch.float32, True),
+            ([-1], torch.float32, True),
+        ]
+    )
+    def forward(self, lhs, rhs):
+        return torch.outer(lhs, rhs)
+
+
+@register_test_case(module_factory=lambda: AtenOuterDynamic())
+def AtenOuterDynamic_basic(module, tu: TestUtils):
+    module.forward(tu.rand(5), tu.rand(5))
+
+
+@register_test_case(module_factory=lambda: AtenOuterDynamic())
+def AtenOuterDynamic_lhs_larger(module, tu: TestUtils):
+    module.forward(tu.rand(7), tu.rand(4))
+
+
+@register_test_case(module_factory=lambda: AtenOuterDynamic())
+def AtenOuterDynamic_rhs_larger(module, tu: TestUtils):
+    module.forward(tu.rand(2), tu.rand(6))