[MLIR][TORCH] Add E2E support for aten.bitwise_and.tensor op

This commit adds lowering of `aten.bitwise_and.tensor` op.

Signed-Off By: Vivek Khandelwal vivek@nod-labs.com

e2e_testing/torchscript/elementwise.py

@@ -552,3 +552,24 @@ class ElementwiseDivScalarModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: ElementwiseDivScalarModule())
 def ElementwiseDivScalarModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
+
+# ==============================================================================
+class ElementwiseAndIntegerModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.int32, True),
+        ([-1, -1], torch.int64, True),
+    ])
+
+    def forward(self, x, y):
+        return torch.bitwise_and(x, y)
+
+
+@register_test_case(module_factory=lambda: ElementwiseAndIntegerModule())
+def ElementwiseAndIntegerModule_basic(module, tu: TestUtils):
+    module.forward(torch.randint(-10, 10, (3, 4)).to(torch.int32),
+                   torch.randint(-10, 10, (3, 4)))

include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td

@@ -1494,6 +1494,21 @@ def Torch_Aten_SoftmaxOp : Torch_Op<"aten._softmax", [
   let assemblyFormat = "$self `,` $dim `,` $half_to_float attr-dict `:` type($self) `,` type($dim) `,` type($half_to_float) `->` type($result)";
 }
 
+def Torch_AtenBitwiseAndTensorOp : Torch_Op<"aten.bitwise_and.Tensor", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::bitwise_and.Tensor : (Tensor, Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchTensorType:$other
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self `,` $other attr-dict `:` type($self) `,` type($other) `->` type($result)";
+}
+
 def Torch_AtenUnsqueezeOp : Torch_Op<"aten.unsqueeze", [
     AllowsTypeRefinement
   ]> {

lib/Conversion/TorchToLinalg/TorchToLinalg.cpp

@@ -1394,6 +1394,22 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
     return b.create<math::SqrtOp>(loc, payloadArgs[0]);
   if (isa<AtenRsqrtOp>(op))
     return b.create<math::RsqrtOp>(loc, payloadArgs[0]);
+  if (auto bitwiseAndTensor = dyn_cast<AtenBitwiseAndTensorOp>(op)) {
+    if (bitwiseAndTensor.getType()
+            .cast<ValueTensorType>()
+            .getDtype()
+            .isa<mlir::FloatType>()) {
+      bitwiseAndTensor.emitError(
+          "Bitwise_And does not support floating point dtype");
+      return nullptr;
+    }
+    Type dtype = converter->convertType(bitwiseAndTensor.getType())
+                     .cast<RankedTensorType>()
+                     .getElementType();
+    Value lhs = convertScalarToDtype(b, loc, payloadArgs[0], dtype);
+    Value rhs = convertScalarToDtype(b, loc, payloadArgs[1], dtype);
+    return b.create<arith::AndIOp>(loc, lhs, rhs);
+  }
   if (isa<AtenLog2Op>(op))
     return b.create<math::Log2Op>(loc, payloadArgs[0]);
   if (isa<AtenAbsOp>(op))
@@ -1895,13 +1911,14 @@ struct ConvertElementwiseOp : ConversionPattern {
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    if (!isa<AtenTanhOp, AtenReluOp, AtenLeakyReluOp, AtenGeluOp, AtenGeluBackwardOp,
-             AtenAddTensorOp, AtenMulTensorOp, AtenDivTensorOp, AtenSubTensorOp,
-             AtenLerpTensorOp, AtenSigmoidOp, AtenExpOp, AtenMinimumOp,
-             AtenMaximumOp, AtenToDtypeOp, AtenClampOp, AtenRsubScalarOp,
-             AtenMulScalarOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
-             AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp, AtenDivScalarOp,
-             AtenAbsOp, AtenReciprocalOp>(op))
+    if (!isa<AtenTanhOp, AtenReluOp, AtenLeakyReluOp, AtenGeluOp,
+             AtenGeluBackwardOp, AtenAddTensorOp, AtenMulTensorOp,
+             AtenDivTensorOp, AtenSubTensorOp, AtenLerpTensorOp, AtenSigmoidOp,
+             AtenExpOp, AtenMinimumOp, AtenMaximumOp, AtenToDtypeOp,
+             AtenClampOp, AtenRsubScalarOp, AtenMulScalarOp, AtenLogOp,
+             AtenSqrtOp, AtenFloorOp, AtenPowTensorScalarOp, AtenLog2Op,
+             AtenRsqrtOp, AtenDivScalarOp, AtenAbsOp, AtenReciprocalOp,
+             AtenBitwiseAndTensorOp>(op))
       return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
 
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
@@ -3193,12 +3210,12 @@ public:
     target.addIllegalOp<AtenBatchNormOp>();
     patterns.add<ConvertAtenBatchNormOp>(typeConverter, context);
     target.addIllegalOp<
-        AtenTanhOp, AtenReluOp, AtenLeakyReluOp, AtenGeluOp, AtenGeluBackwardOp, AtenAddTensorOp,
-        AtenMulTensorOp, AtenDivTensorOp, AtenSubTensorOp, AtenLerpTensorOp,
-        AtenSigmoidOp, AtenMinimumOp, AtenMaximumOp, AtenToDtypeOp, AtenClampOp,
-        AtenRsubScalarOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
-        AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp, AtenAbsOp,
-        AtenReciprocalOp>();
+        AtenTanhOp, AtenReluOp, AtenLeakyReluOp, AtenGeluOp, AtenGeluBackwardOp,
+        AtenAddTensorOp, AtenMulTensorOp, AtenDivTensorOp, AtenSubTensorOp,
+        AtenLerpTensorOp, AtenSigmoidOp, AtenMinimumOp, AtenMaximumOp,
+        AtenToDtypeOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp, AtenSqrtOp,
+        AtenFloorOp, AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp, AtenAbsOp,
+        AtenReciprocalOp, AtenBitwiseAndTensorOp>();
     patterns.add<ConvertElementwiseOp>(typeConverter, context);
     target.addIllegalOp<AtenSqueezeOp>();
     patterns.add<ConvertAtenSqueezeOp>(typeConverter, context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -304,7 +304,7 @@ public:
       return visitBinaryTensorScalarOp(op, operands);
     } else if (isa<AtenAddTensorOp, AtenSubTensorOp, AtenMulTensorOp,
                    AtenDivTensorOp, Aten__And__TensorOp, AtenEqTensorOp,
-                   AtenMinimumOp, AtenMaximumOp>(op)) {
+                   AtenMinimumOp, AtenMaximumOp, AtenBitwiseAndTensorOp>(op)) {
       return visitBinaryBroadcastingOp(op, operands);
     } else if (auto lerpTensor = llvm::dyn_cast<AtenLerpTensorOp>(op)) {
       return visitAtenLerpTensorOp(lerpTensor, operands);

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -470,6 +470,8 @@ def emit_aten_ops(torch_ir_dir: str, registry: Registry):
                 "aten::rsqrt : (Tensor) -> (Tensor)",
                 "aten::abs : (Tensor) -> (Tensor)",
                 "aten::reciprocal : (Tensor) -> (Tensor)",
+                "aten::bitwise_and.Tensor : (Tensor, Tensor) -> (Tensor)",
+
         ]:
             emit_with_mutating_variants(key)
         # Elementwise tensor compute ops that don't have the standard mutating