[TorchToArith] Add a lowering for torch.add.float_int (#3594)

llvm · Aug 7, 2024 · 8d95fe9 · 8d95fe9
1 parent a51b4e0
commit 8d95fe9
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 3 deletions.
diff --git a/lib/Conversion/TorchToArith/TorchToArith.cpp b/lib/Conversion/TorchToArith/TorchToArith.cpp
@@ -72,8 +72,11 @@ class ConvertAtenBinaryOp : public OpConversionPattern<AtenOp> {
   matchAndRewrite(AtenOp op,
                   typename OpConversionPattern<AtenOp>::OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.template replaceOpWithNewOp<BinOp>(op, adaptor.getA(),
-                                                adaptor.getB());
+    Value a = adaptor.getA();
+    Value b = adaptor.getB();
+    if (llvm::is_one_of<AtenOp, AtenAddFloatIntOp>::value)
+      b = convertScalarToDtype(rewriter, op.getLoc(), b, a.getType());
+    rewriter.template replaceOpWithNewOp<BinOp>(op, a, b);
     return success();
   }
 };
@@ -255,6 +258,25 @@ class ConvertAtenCastOp : public OpConversionPattern<AtenOp> {
 };
 } // namespace
 
+namespace {
+template <typename AtenOp>
+class ConvertAtenScalarArithOp : public OpConversionPattern<AtenOp> {
+public:
+  using OpConversionPattern<AtenOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(AtenOp op,
+                  typename OpConversionPattern<AtenOp>::OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type resultType =
+        this->getTypeConverter()->convertType(op->getResult(0).getType());
+    Value result =
+        convertScalarToDtype(rewriter, op.getLoc(), adaptor.getA(), resultType);
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class ConvertAtenAddOp : public OpConversionPattern<AtenAddOp> {
 public:
@@ -444,9 +466,12 @@ class ConvertTorchToArith
     target.addIllegalOp<AtenAddOp>();
     patterns.add<ConvertAtenAddOp>(typeConverter, context);
 
-    target.addIllegalOp<AtenAddIntOp, AtenSubIntOp, AtenMulIntOp>();
+    target.addIllegalOp<AtenAddIntOp, AtenAddFloatIntOp, AtenSubIntOp,
+                        AtenMulIntOp>();
     patterns.add<ConvertAtenBinaryOp<AtenAddIntOp, arith::AddIOp>>(
         typeConverter, context);
+    patterns.add<ConvertAtenBinaryOp<AtenAddFloatIntOp, arith::AddFOp>>(
+        typeConverter, context);
     patterns.add<ConvertAtenBinaryOp<AtenSubIntOp, arith::SubIOp>>(
         typeConverter, context);
     patterns.add<ConvertAtenBinaryOp<AtenMulIntOp, arith::MulIOp>>(

diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -117,6 +117,7 @@
     # END tests failing due to: torch._dynamo.exc.Unsupported: call_function BuiltinVariable(float) [TensorVariable()] {}
     # START tests failing due to: torch._dynamo.exc.Unsupported: call_function BuiltinVariable(int) [TensorVariable()] {}
     "AddIntModule_basic",
+    "AddFloatIntModule_basic",
     "AtenIntTensorCharDtypeModule_basic",
     "BoolIntFalseModule_basic",
     "BoolIntTrueModule_basic",
@@ -339,6 +340,7 @@
 
 FX_IMPORTER_XFAIL_SET = {
     "ReduceAnyDimFloatModule_basic",
+    "AddFloatIntModule_basic",
     "AllBoolFalseModule_basic",
     "AllBoolTrueModule_basic",
     "AnyBoolFalseModule_basic",
@@ -855,6 +857,7 @@
     "AdaptiveAvgPool2dOutputSizeDivisibleByInputStaticModule_basic",
     "AdaptiveAvgPool2dFixedKernelStrideSizeStaticModule_basic",
     "AddIntModule_basic",
+    "AddFloatIntModule_basic",
     "AliasModule_basic",
     "TrueFalseOrBoolOpModule_basic",
     "AllBoolFalseModule_basic",
@@ -2100,6 +2103,7 @@
     "_ConvolutionDeprecated2DDeterministicModule_basic",
     "MaxPool3dEmptyStrideStaticModule_basic",
     "AddIntModule_basic",
+    "AddFloatIntModule_basic",
     "ArangeStartOutViewModule_basic",
     "AtenIntBoolOpModule_basic",
     "BernoulliTensorModule_basic",
@@ -2288,6 +2292,7 @@
     "AdaptiveMaxPool3dStatic_basic",
     "AddCDivModule_basic",
     "AddIntModule_basic",
+    "AddFloatIntModule_basic",
     "Add_Module_basic",
     "AllBoolFalseModule_basic",
     "AllBoolTrueModule_basic",
@@ -2840,6 +2845,7 @@
     "AdaptiveMaxPool3dStaticWithIndices_basic",
     "AdaptiveMaxPool3dStatic_basic",
     "AddIntModule_basic",
+    "AddFloatIntModule_basic",
     "Add_MixPModule_basic",
     "AllBoolFalseModule_basic",
     "AllBoolTrueModule_basic",
@@ -3609,6 +3615,7 @@
     "AdaptiveMaxPool3dStatic_basic",
     "AddCDivModule_basic",
     "AddIntModule_basic",
+    "AddFloatIntModule_basic",
     "AddSizeIntModule_basic",
     "AddSizeIntNegDimModule_basic",
     "Add_MixPModule_basic",

diff --git a/projects/pt1/python/torch_mlir_e2e_test/test_suite/scalar.py b/projects/pt1/python/torch_mlir_e2e_test/test_suite/scalar.py
@@ -36,6 +36,30 @@ def AddIntModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class AddFloatIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([], torch.float32, True),
+            ([], torch.int64, True),
+        ]
+    )
+    def forward(self, lhs, rhs):
+        return float(lhs) + int(rhs)
+
+
+@register_test_case(module_factory=lambda: AddFloatIntModule())
+def AddFloatIntModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(), tu.randint(low=-100, high=100))
+
+
+# ==============================================================================
+
+
 class SubIntModule(torch.nn.Module):
     def __init__(self):
         super().__init__()