[TorchToArith] Implement conversion patterns for AtenSubOp, AtenMulOp, AtenDivOp.

Implement conversion patterns for `AtenSubOp`, `AtenMulOp`, `AtenDivOp`:
Use unified template patterns `ConvertAtenBinaryScalarOp` to handle `AtenAddOp`, `AtenSubOp`, `AtenMulOp`;
Use unified template patterns `ConvertAtenDivOp` to handle `AtenDivOp`, `AtenDivIntOp`, `AtenDivFloatOp`.

Author: Junru-Sheng

lib/Conversion/TorchToArith/TorchToArith.cpp

@@ -134,12 +134,13 @@ class ConvertAtenUnaryOpToFloatMathOp : public OpConversionPattern<AtenOp> {
 } // namespace
 
 namespace {
-class ConvertAtenDivIntOp : public OpConversionPattern<AtenDivIntOp> {
+template <typename AtenOp>
+class ConvertAtenDivOp : public OpConversionPattern<AtenOp> {
 public:
-  using OpConversionPattern<AtenDivIntOp>::OpConversionPattern;
+  using OpConversionPattern<AtenOp>::OpConversionPattern;
+  using OpAdaptor = typename AtenOp::Adaptor;
   LogicalResult
-  matchAndRewrite(AtenDivIntOp op,
-                  typename OpConversionPattern<AtenDivIntOp>::OpAdaptor adaptor,
+  matchAndRewrite(AtenOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     Value a = convertScalarToDtype(rewriter, loc, adaptor.getA(),
@@ -306,11 +307,13 @@ class ConvertAtenScalarArithOp : public OpConversionPattern<AtenOp> {
 } // namespace
 
 namespace {
-class ConvertAtenAddOp : public OpConversionPattern<AtenAddOp> {
+template <typename AtenOp, typename ArithFOp, typename ArithIOp>
+class ConvertAtenBinaryScalarOp : public OpConversionPattern<AtenOp> {
 public:
-  using OpConversionPattern::OpConversionPattern;
+  using OpConversionPattern<AtenOp>::OpConversionPattern;
+  using OpAdaptor = typename AtenOp::Adaptor;
   LogicalResult
-  matchAndRewrite(AtenAddOp op, OpAdaptor adaptor,
+  matchAndRewrite(AtenOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     Type resultType =
@@ -320,9 +323,9 @@ class ConvertAtenAddOp : public OpConversionPattern<AtenAddOp> {
     Value operandB =
         convertScalarToDtype(rewriter, loc, adaptor.getB(), resultType);
     if (isa<mlir::FloatType>(resultType)) {
-      rewriter.replaceOpWithNewOp<arith::AddFOp>(op, operandA, operandB);
+      rewriter.replaceOpWithNewOp<ArithFOp>(op, operandA, operandB);
     } else if (isa<mlir::IntegerType>(resultType)) {
-      rewriter.replaceOpWithNewOp<arith::AddIOp>(op, operandA, operandB);
+      rewriter.replaceOpWithNewOp<ArithIOp>(op, operandA, operandB);
     } else {
       return rewriter.notifyMatchFailure(
           op, "unimplemented: only support integer or float result type");
@@ -497,8 +500,17 @@ class ConvertTorchToArith
     patterns.add<ConvertAtenCastOp<AtenFloatScalarOp>>(typeConverter, context);
     patterns.add<ConvertAtenCastOp<AtenIntScalarOp>>(typeConverter, context);
 
-    target.addIllegalOp<AtenAddOp>();
-    patterns.add<ConvertAtenAddOp>(typeConverter, context);
+    target.addIllegalOp<AtenAddOp, AtenSubOp, AtenMulOp>();
+    patterns.add<
+        ConvertAtenBinaryScalarOp<AtenAddOp, airth::AddFOp, airth::AddIOp>>(
+        typeConverter, context);
+    patterns.add<
+        ConvertAtenBinaryScalarOp<AtenSubOp, airth::SubFOp, airth::SubIOp>>(
+        typeConverter, context);
+    patterns.add<
+        ConvertAtenBinaryScalarOp<AtenMulOp, airth::MulFOp, airth::MulIOp>>(
+        typeConverter, context);
+
     target.addIllegalOp<AtenNegIntOp>();
     patterns.add<ConvertAtenNegIntOp>(typeConverter, context);
     target.addIllegalOp<AtenAddIntOp, AtenAddFloatIntOp, AtenSubIntOp,
@@ -523,11 +535,12 @@ class ConvertTorchToArith
         typeConverter, context);
     patterns.add<ConvertAtenBinaryOp<AtenMulFloatOp, arith::MulFOp>>(
         typeConverter, context);
-    target.addIllegalOp<AtenDivIntOp>();
-    patterns.add<ConvertAtenDivIntOp>(typeConverter, context);
-    target.addIllegalOp<AtenDivFloatOp>();
-    patterns.add<ConvertAtenBinaryOp<AtenDivFloatOp, arith::DivFOp>>(
-        typeConverter, context);
+
+    target.addIllegalOp<AtenDivOp, AtenDivIntOp, AtenDivFloatOp>();
+    patterns.add<ConvertAtenDivOp<AtenDivOp>>(typeConverter, context);
+    patterns.add<ConvertAtenDivOp<AtenDivIntOp>>(typeConverter, context);
+    patterns.add<ConvertAtenDivOp<AtenDivFloatOp>>(typeConverter, context);
+
     target.addIllegalOp<AtenFloordivIntOp>();
     patterns.add<ConvertAtenBinaryOp<AtenFloordivIntOp, arith::FloorDivSIOp>>(
         typeConverter, context);