[LINALG] Add E2E support for `aten.eq.int` op

- This commit adds lowering of `aten.eq.int` op as a part of
  `convert-torch-to-std` pass.
- It also refactors the code for binary comparison ops lowering.

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

e2e_testing/torchscript/main.py

@@ -41,6 +41,7 @@ from . import argmax
 from . import matmul
 from . import reshape_like
 from . import scalar
+from . import scalar_comparison
 from . import squeeze
 from . import slice_like
 from . import nll_loss

e2e_testing/torchscript/scalar_comparison.py

@@ -0,0 +1,71 @@
+# 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 NeIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([], torch.int64, True),
+        ([], torch.int64, True),
+    ])
+    def forward(self, lhs, rhs):
+        return int(lhs) != int(rhs)
+
+
+@register_test_case(module_factory=lambda: NeIntModule())
+def NeIntModule_basic(module, tu: TestUtils):
+    module.forward(torch.randint(-100, 100, ()), torch.randint(-100, 100, ()))
+
+# ==============================================================================
+
+class EqIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([], torch.int64, True),
+        ([], torch.int64, True),
+    ])
+    def forward(self, lhs, rhs):
+        return int(lhs) == int(rhs)
+
+
+@register_test_case(module_factory=lambda: EqIntModule())
+def EqIntModule_basic(module, tu: TestUtils):
+    module.forward(torch.randint(-100, 100, ()), torch.randint(-100, 100, ()))
+
+# ==============================================================================
+
+class GtIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([], torch.int64, True),
+        ([], torch.int64, True),
+    ])
+    def forward(self, lhs, rhs):
+        return int(lhs) > int(rhs)
+
+
+@register_test_case(module_factory=lambda: GtIntModule())
+def GtIntModule_basic(module, tu: TestUtils):
+    module.forward(torch.randint(-100, 100, ()), torch.randint(-100, 100, ()))
+

lib/Conversion/TorchToStd/TorchToStd.cpp

@@ -75,28 +75,17 @@ public:
 } // namespace
 
 namespace {
-class ConvertAtenNeIntOp : public OpConversionPattern<AtenNeIntOp> {
+// Lowers aten integer comparison ops.
+template <typename AtenOp, arith::CmpIPredicate Pred>
+class ConvertAtenIntComparisonOp : public OpConversionPattern<AtenOp> {
 public:
-  using OpConversionPattern::OpConversionPattern;
+  using OpConversionPattern<AtenOp>::OpConversionPattern;
   LogicalResult
-  matchAndRewrite(AtenNeIntOp op, OpAdaptor adaptor,
+  matchAndRewrite(AtenOp op,
+                  typename OpConversionPattern<AtenOp>::OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<arith::CmpIOp>(op, arith::CmpIPredicate::ne,
-                                               adaptor.a(), adaptor.b());
-    return success();
-  }
-};
-} // namespace
-
-namespace {
-class ConvertAtenGtIntOp : public OpConversionPattern<AtenGtIntOp> {
-public:
-  using OpConversionPattern::OpConversionPattern;
-  LogicalResult
-  matchAndRewrite(AtenGtIntOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<arith::CmpIOp>(op, arith::CmpIPredicate::sgt,
-                                               adaptor.a(), adaptor.b());
+    rewriter.replaceOpWithNewOp<arith::CmpIOp>(op, Pred, adaptor.a(),
+                                               adaptor.b());
     return success();
   }
 };
@@ -189,10 +178,16 @@ public:
     patterns.add<ConvertAtenDimOp>(typeConverter, context);
     target.addIllegalOp<RuntimeAssertOp>();
     patterns.add<ConvertRuntimeAssertOp>(typeConverter, context);
-    target.addIllegalOp<AtenNeIntOp>();
-    patterns.add<ConvertAtenNeIntOp>(typeConverter, context);
-    target.addIllegalOp<AtenGtIntOp>();
-    patterns.add<ConvertAtenGtIntOp>(typeConverter, context);
+    target.addIllegalOp<AtenNeIntOp, AtenEqIntOp, AtenGtIntOp>();
+    patterns
+        .add<ConvertAtenIntComparisonOp<AtenNeIntOp, arith::CmpIPredicate::ne>>(
+            typeConverter, context);
+    patterns
+        .add<ConvertAtenIntComparisonOp<AtenEqIntOp, arith::CmpIPredicate::eq>>(
+            typeConverter, context);
+    patterns.add<
+        ConvertAtenIntComparisonOp<AtenGtIntOp, arith::CmpIPredicate::sgt>>(
+        typeConverter, context);
     target.addIllegalOp<ValueTensorLiteralOp>();
     patterns.add<ConvertTorchTensorLiteralOp>(typeConverter, context);
 

test/Conversion/TorchToStd/basic.mlir

@@ -40,6 +40,19 @@ func @torch.aten.ne.int(%arg0: !torch.int, %arg1: !torch.int) -> !torch.bool {
   return %0 : !torch.bool
 }
 
+// CHECK-LABEL:   func @torch.aten.eq.int(
+// CHECK-SAME:                            %[[LHS:.*]]: !torch.int,
+// CHECK-SAME:                            %[[RHS:.*]]: !torch.int) -> !torch.bool {
+// CHECK:           %[[LHS_I64:.*]] = torch_c.to_i64 %[[LHS]]
+// CHECK:           %[[RHS_I64:.*]] = torch_c.to_i64 %[[RHS]]
+// CHECK:           %[[CMP:.*]] = arith.cmpi eq, %[[LHS_I64]], %[[RHS_I64]] : i64
+// CHECK:           %[[CMP_TORCH_BOOL:.*]] = torch_c.from_i1 %[[CMP]]
+// CHECK:           return %[[CMP_TORCH_BOOL]] : !torch.bool
+func @torch.aten.eq.int(%arg0: !torch.int, %arg1: !torch.int) -> !torch.bool {
+  %0 = torch.aten.eq.int %arg0, %arg1 : !torch.int, !torch.int -> !torch.bool
+  return %0 : !torch.bool
+}
+
 // CHECK-LABEL:   func @torch.aten.gt.int(
 // CHECK-SAME:                            %[[LHS:.*]]: !torch.int,
 // CHECK-SAME:                            %[[RHS:.*]]: !torch.int) -> !torch.bool {