[MLIR][TORCH] Add E2E support for `aten.[div.int|bitwise_or.Tensor]` ops

This commit adds lowering of `aten.div.int` and `aten.bitwise_or.Tensor`
ops. Both these ops are required in order to support bloom_560m model.

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

e2e_testing/xfail_sets.py

@@ -574,6 +574,7 @@ LTC_XFAIL_SET = {
     "LiftFreshCopyModule_basic",
     "Matmul_dot",
     "MulIntModule_basic",
+    "DivIntModule_basic",
     "NeFloatIntModule_basic",
     "NeIntModule_basic",
     "NewEmptyModuleDefaultDtype_basic",

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

@@ -2192,6 +2192,53 @@ def Torch_AtenBitwiseAnd_TensorOp : Torch_Op<"aten.bitwise_and_.Tensor", [
   }];
 }
 
+def Torch_AtenBitwiseOrTensorOp : Torch_Op<"aten.bitwise_or.Tensor", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::bitwise_or.Tensor : (Tensor, Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchTensorType:$other
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenBitwiseOrTensorOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenBitwiseOrTensorOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+}
+
+def Torch_AtenBitwiseOr_TensorOp : Torch_Op<"aten.bitwise_or_.Tensor", [
+    IsTrailingUnderscoreInplaceVariant,
+    AllowsTypeRefinement
+  ]> {
+  let summary = "Generated op for `aten::bitwise_or_.Tensor : (Tensor, Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchTensorType:$other
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenBitwiseOr_TensorOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenBitwiseOr_TensorOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+}
+
 def Torch_AtenThresholdOp : Torch_Op<"aten.threshold", [
     AllowsTypeRefinement,
     HasValueSemantics,
@@ -8405,6 +8452,31 @@ def Torch_AtenMulIntOp : Torch_Op<"aten.mul.int", [
   let hasFolder = 1;
 }
 
+def Torch_AtenDivIntOp : Torch_Op<"aten.div.int", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::div.int : (int, int) -> (float)`";
+  let arguments = (ins
+    Torch_IntType:$a,
+    Torch_IntType:$b
+  );
+  let results = (outs
+    Torch_FloatType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenDivIntOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenDivIntOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+  let hasFolder = 1;
+}
+
 def Torch_AtenNegIntOp : Torch_Op<"aten.neg.int", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Conversion/TorchToArith/TorchToArith.cpp

@@ -121,6 +121,25 @@ public:
 };
 } // namespace
 
+namespace {
+class ConvertAtenDivIntOp : public OpConversionPattern<AtenDivIntOp> {
+public:
+  using OpConversionPattern<AtenDivIntOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(AtenDivIntOp op,
+                  typename OpConversionPattern<AtenDivIntOp>::OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    Value a =
+        convertScalarToDtype(rewriter, loc, adaptor.a(), rewriter.getF64Type());
+    Value b =
+        convertScalarToDtype(rewriter, loc, adaptor.b(), rewriter.getF64Type());
+    rewriter.replaceOpWithNewOp<arith::DivFOp>(op, a, b);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 // Lowers aten integer comparison ops.
 template <typename AtenOp, arith::CmpIPredicate Pred>
@@ -374,6 +393,8 @@ public:
     target.addIllegalOp<AtenSubFloatOp>();
     patterns.add<ConvertAtenBinaryOp<AtenSubFloatOp, arith::SubFOp>>(
         typeConverter, context);
+    target.addIllegalOp<AtenDivIntOp>();
+    patterns.add<ConvertAtenDivIntOp>(typeConverter, context);
     target.addIllegalOp<AtenDivFloatOp>();
     patterns.add<ConvertAtenBinaryOp<AtenDivFloatOp, arith::DivFOp>>(
         typeConverter, context);

lib/Conversion/TorchToLinalg/Uncategorized.cpp

@@ -199,6 +199,22 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
     Value rhs = convertScalarToDtype(b, loc, payloadArgs[1], dtype);
     return b.create<arith::AndIOp>(loc, lhs, rhs);
   }
+  if (auto bitwiseOrTensor = dyn_cast<AtenBitwiseOrTensorOp>(op)) {
+    if (bitwiseOrTensor.getType()
+            .cast<ValueTensorType>()
+            .getDtype()
+            .isa<mlir::FloatType>()) {
+      bitwiseOrTensor.emitError(
+          "Bitwise_Or does not support floating point dtype");
+      return nullptr;
+    }
+    Type dtype = converter->convertType(bitwiseOrTensor.getType())
+                     .cast<RankedTensorType>()
+                     .getElementType();
+    Value lhs = convertScalarToDtype(b, loc, payloadArgs[0], dtype);
+    Value rhs = convertScalarToDtype(b, loc, payloadArgs[1], dtype);
+    return b.create<arith::OrIOp>(loc, lhs, rhs);
+  }
   if (auto logicalOr = dyn_cast<AtenLogicalOrOp>(op)) {
     MLIRContext *context = op->getContext();
     Type floatDtype = mlir::FloatType::getF64(context);
@@ -1006,12 +1022,12 @@ public:
              AtenSqrtOp, AtenFloorOp, AtenPowTensorScalarOp,
              AtenPowTensorTensorOp, AtenLog2Op, AtenLog1pOp, AtenRsqrtOp,
              AtenDivScalarOp, AtenRemainderScalarOp, AtenAbsOp,
-             AtenReciprocalOp, AtenBitwiseAndTensorOp, AtenGtScalarOp,
-             AtenGeScalarOp, AtenEqScalarOp, AtenLtScalarOp, AtenLeScalarOp,
-             AtenWhereSelfOp, AtenCeilOp, AtenGtTensorOp, AtenEqTensorOp,
-             AtenLtTensorOp, AtenSubScalarOp, AtenAddScalarOp, AtenThresholdOp,
-             AtenThresholdBackwardOp, AtenCloneOp, AtenSinOp, AtenCosOp,
-             AtenNeScalarOp, AtenNegOp, AtenMaskedFillScalarOp,
+             AtenReciprocalOp, AtenBitwiseAndTensorOp, AtenBitwiseOrTensorOp,
+             AtenGtScalarOp, AtenGeScalarOp, AtenEqScalarOp, AtenLtScalarOp,
+             AtenLeScalarOp, AtenWhereSelfOp, AtenCeilOp, AtenGtTensorOp,
+             AtenEqTensorOp, AtenLtTensorOp, AtenSubScalarOp, AtenAddScalarOp,
+             AtenThresholdOp, AtenThresholdBackwardOp, AtenCloneOp, AtenSinOp,
+             AtenCosOp, AtenNeScalarOp, AtenNegOp, AtenMaskedFillScalarOp,
              AtenMaskedFillTensorOp, AtenLogicalOrOp, AtenTriuOp,
              AtenBitwiseNotOp>(op))
       return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
@@ -1483,10 +1499,10 @@ void mlir::torch::torch_to_linalg::populateUncategorizedPatternsAndLegality(
       AtenLogOp, AtenErfOp, AtenSqrtOp, AtenFloorOp, AtenCeilOp,
       AtenPowTensorScalarOp, AtenPowTensorTensorOp, AtenLog2Op, AtenLog1pOp,
       AtenRsqrtOp, AtenAbsOp, AtenReciprocalOp, AtenBitwiseAndTensorOp,
-      AtenGtScalarOp, AtenGeScalarOp, AtenEqScalarOp, AtenLtScalarOp,
-      AtenLeScalarOp, AtenWhereSelfOp, AtenGtTensorOp, AtenEqTensorOp,
-      AtenLtTensorOp, AtenThresholdOp, AtenThresholdBackwardOp, AtenCloneOp,
-      AtenSinOp, AtenCosOp, AtenNeScalarOp, AtenMaskedFillScalarOp,
+      AtenBitwiseOrTensorOp, AtenGtScalarOp, AtenGeScalarOp, AtenEqScalarOp,
+      AtenLtScalarOp, AtenLeScalarOp, AtenWhereSelfOp, AtenGtTensorOp,
+      AtenEqTensorOp, AtenLtTensorOp, AtenThresholdOp, AtenThresholdBackwardOp,
+      AtenCloneOp, AtenSinOp, AtenCosOp, AtenNeScalarOp, AtenMaskedFillScalarOp,
       AtenMaskedFillTensorOp, AtenLogicalOrOp, AtenTriuOp,
       AtenRemainderScalarOp, AtenBitwiseNotOp>();
   patterns.add<ConvertElementwiseOp>(typeConverter, context);

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -2175,6 +2175,20 @@ OpFoldResult AtenDivFloatOp::fold(ArrayRef<Attribute> operands) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// AtenDivIntOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AtenDivIntOp::fold(ArrayRef<Attribute> operands) {
+  int64_t lhs, rhs;
+  bool lConstant = matchPattern(getOperand(0), m_TorchConstantInt(&lhs));
+  bool rConstant = matchPattern(getOperand(1), m_TorchConstantInt(&rhs));
+  if (lConstant && rConstant)
+    return getF64FloatAttr(getContext(), double(lhs) / rhs);
+  return nullptr;
+}
+
+//===----------------------------------------------------------------------===//
 // AtenCeilFloatOp
 //===----------------------------------------------------------------------===//
 
@@ -2185,8 +2199,6 @@ OpFoldResult AtenCeilFloatOp::fold(ArrayRef<Attribute> operands) {
   return nullptr;
 }
 
-//===----------------------------------------------------------------------===//
-
 //===----------------------------------------------------------------------===//
 // PrimMaxIntOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -767,8 +767,8 @@ void TypeAnalysis::visitOperation(Operation *op,
   // Promote the two dtypes assuming possibly-zero rank.
   if (isa<AtenAddTensorOp, AtenSubTensorOp, AtenMulTensorOp, AtenDivTensorOp,
           AtenDivTensorModeOp, Aten__And__TensorOp, AtenMinimumOp,
-          AtenMaximumOp, AtenBitwiseAndTensorOp, AtenThresholdBackwardOp,
-          AtenFloorDivideOp>(op)) {
+          AtenMaximumOp, AtenBitwiseAndTensorOp, AtenBitwiseOrTensorOp,
+          AtenThresholdBackwardOp, AtenFloorDivideOp>(op)) {
     auto knowledge =
         ValueKnowledge::getTensorPessimisticValueState(op->getContext());
     knowledge.dtype = getPromotedResultType(

lib/Dialect/Torch/Transforms/ShapeLibrary.cpp

@@ -6383,6 +6383,10 @@ StringRef mlir::torch::Torch::getShapeLibrary() {
 "    %0 = call @__torch__.torch.jit._shape_functions.broadcast(%arg0, %arg1) : (!torch.list<int>, !torch.list<int>) -> !torch.list<int>\n"
 "    return %0 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.bitwise_or.Tensor\"(%arg0: !torch.list<int>, %arg1: !torch.list<int>) -> !torch.list<int> {\n"
+"    %0 = call @__torch__.torch.jit._shape_functions.broadcast(%arg0, %arg1) : (!torch.list<int>, !torch.list<int>) -> !torch.list<int>\n"
+"    return %0 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.bitwise_and.Tensor\"(%arg0: !torch.list<int>, %arg1: !torch.list<int>) -> !torch.list<int> {\n"
 "    %0 = call @__torch__.torch.jit._shape_functions.broadcast(%arg0, %arg1) : (!torch.list<int>, !torch.list<int>) -> !torch.list<int>\n"
 "    return %0 : !torch.list<int>\n"

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

@@ -866,6 +866,9 @@ def aten〇minimum(self: List[int], other: List[int]) -> List[int]:
 def aten〇maximum(self: List[int], other: List[int]) -> List[int]:
     return upstream_shape_functions.broadcast(self, other)
 
+def aten〇bitwise_or〇Tensor(self: List[int], other: List[int]) -> List[int]:
+    return upstream_shape_functions.broadcast(self, other)
+
 def aten〇bitwise_and〇Tensor(self: List[int], other: List[int]) -> List[int]:
     return upstream_shape_functions.broadcast(self, other)
 

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

@@ -287,6 +287,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
             "aten::abs : (Tensor) -> (Tensor)",
             "aten::reciprocal : (Tensor) -> (Tensor)",
             "aten::bitwise_and.Tensor : (Tensor, Tensor) -> (Tensor)",
+            "aten::bitwise_or.Tensor : (Tensor, Tensor) -> (Tensor)",
             "aten::threshold : (Tensor, Scalar, Scalar) -> (Tensor)",
             "aten::square : (Tensor) -> (Tensor)",
             "aten::unsqueeze : (Tensor, int) -> (Tensor)",
@@ -570,6 +571,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
     emit("aten::add.int : (int, int) -> (int)", has_folder=True)
     emit("aten::sub.int : (int, int) -> (int)", has_folder=True)
     emit("aten::mul.int : (int, int) -> (int)", has_folder=True)
+    emit("aten::div.int : (int, int) -> (float)", has_folder=True)
     emit("aten::neg.int : (int) -> (int)", has_folder=True)
     emit("aten::log.int : (int) -> (float)")
     emit("aten::add.float_int : (float, int) -> (float)")

python/torch_mlir_e2e_test/test_suite/elementwise.py

@@ -1553,6 +1553,31 @@ def ElementwiseAndIntegerModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class ElementwiseOrIntegerModule(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_or(x, y)
+
+
+@register_test_case(module_factory=lambda: ElementwiseOrIntegerModule())
+def ElementwiseOrIntegerModule_basic(module, tu: TestUtils):
+    module.forward(
+        tu.randint(3, 4, low=-10, high=10).to(torch.int32),
+        tu.randint(3, 4, low=-10, high=10))
+
+
+# ==============================================================================
+
+
 class ElementwiseNotIntegerModule(torch.nn.Module):
 
     def __init__(self):

python/torch_mlir_e2e_test/test_suite/scalar.py

@@ -104,6 +104,31 @@ def MulIntModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class DivIntModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([], torch.int64, True),
+        ([], torch.int64, True),
+    ])
+    def forward(self, lhs, rhs):
+        # Cast the result to float to make e2e test baseline result to be a float.
+        # Without the cast, baseline result is a Tensor which is unexpected.
+        return float(torch.ops.aten.div(int(lhs), int(rhs)))
+
+
+@register_test_case(module_factory=lambda: DivIntModule())
+def DivIntModule_basic(module, tu: TestUtils):
+    module.forward(tu.randint(low=-10, high=10), tu.randint(low=3, high=10))
+
+
+# ==============================================================================
+
+
 class DivFloatModule(torch.nn.Module):
 
     def __init__(self):

test/Dialect/Torch/canonicalize.mlir

@@ -1351,6 +1351,16 @@ func.func @torch.aten.div.float$fold_cst_operands() -> !torch.float {
   return %0 : !torch.float
 }
 
+// CHECK-LABEL:   func.func @torch.aten.div.int$fold_cst_operands(
+// CHECK:           %[[CST:.*]] = torch.constant.float 5.000000e-01
+// CHECK:           return %[[CST]] : !torch.float
+func.func @torch.aten.div.int$fold_cst_operands() -> !torch.float {
+  %int2 = torch.constant.int 2
+  %int4 = torch.constant.int 4
+  %0 = torch.aten.div.int %int2, %int4 : !torch.int, !torch.int -> !torch.float
+  return %0 : !torch.float
+}
+
 // CHECK-LABEL:   func.func @torch.aten.to.dtype_layout$same_dtype(
 // CHECK-SAME:            %[[ARG:.*]]: !torch.tensor<[?,?],f32>) -> !torch.tensor<[?,?],f32> {
 // CHECK-NEXT:      return %[[ARG]] : !torch.tensor<[?,?],f32>