[MHLO] fix tensor mode aten.div op pattern

Co-authored-by: Bairen Yi <yibairen.byron@bytedance.com>
Co-authored-by: Jiawei Wu <xremold@gmail.com>
Co-authored-by: Tianyou Guo <tianyou.gty@alibaba-inc.com>
Co-authored-by: Xu Yan <yancey.yx@alibaba-inc.com>
Co-authored-by: Ziheng Jiang <ziheng.jiang@bytedance.com>

lib/Conversion/TorchToMhlo/Basic.cpp

@@ -208,7 +208,6 @@ class ConvertAtenMulDivOp : public OpConversionPattern<AtenOpT> {
           "only floating-point or integer datatype legalization supported");
     }
 
-    Value lhsTensor = lhs;
     if (std::is_same<AtenOpT, AtenSquareOp>()) {
       rhs = lhs;
     } else if (!rhsType) {
@@ -217,8 +216,37 @@ class ConvertAtenMulDivOp : public OpConversionPattern<AtenOpT> {
     DenseIntElementsAttr bcastDimensions;
     lhs = mhlo::promoteType(rewriter, lhs, outType);
     rhs = mhlo::promoteType(rewriter, rhs, outType);
-    rewriter.replaceOpWithNewOp<ChloOpT>(op, outType, lhs, rhs,
-                                         bcastDimensions);
+    auto loc = op.getLoc();
+    Value result =
+        rewriter.create<ChloOpT>(loc, outType, lhs, rhs, bcastDimensions);
+
+    if (!isa<AtenDivTensorModeOp>(op)) {
+      rewriter.replaceOp(op, result);
+      return success();
+    }
+
+    AtenDivTensorModeOp divTensorModeOp =
+        llvm::dyn_cast<AtenDivTensorModeOp>(op.getOperation());
+    std::string roundingMode;
+    if (!matchPattern(divTensorModeOp.rounding_mode(),
+                      m_TorchConstantStr(roundingMode)))
+      return rewriter.notifyMatchFailure(
+          op, "only support constant str rounding mode");
+
+    if (roundingMode == "trunc") {
+      // "trunc" - rounds the results of the division towards zero. Equivalent
+      // to C-style integer division.
+      auto sign = rewriter.create<mhlo::SignOp>(loc, result);
+      auto abs = rewriter.create<mhlo::AbsOp>(loc, result);
+      auto floor = rewriter.create<mhlo::FloorOp>(loc, abs);
+      result = rewriter.create<mhlo::MulOp>(loc, sign, floor).getResult();
+    }
+    if (roundingMode == "floor") {
+      // "floor" - rounds the results of the division down. Equivalent to
+      // floor division in Python (the // operator)
+      result = rewriter.create<mhlo::FloorOp>(loc, result).getResult();
+    }
+    rewriter.replaceOp(op, result);
     return success();
   }
 };
@@ -554,7 +582,6 @@ LogicalResult ConvertAtenOp<PrimNumToTensorScalarOp>::matchAndRewrite(
   RankedTensorType outputType = getTypeConverter()
                                     ->convertType(op->getResult(0).getType())
                                     .cast<RankedTensorType>();
-  auto outputShape = outputType.getShape();
   auto outputElemType = outputType.getElementType();
   Value mhloTensor =
       mhlo::scalarToMhloTensor(rewriter, op, adaptor.a(), outputElemType);

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -2167,8 +2167,11 @@ class DecomposeAtenFloorDivideOp : public OpRewritePattern<AtenFloorDivideOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(AtenFloorDivideOp op,
                                 PatternRewriter &rewriter) const override {
+    // https://pytorch.org/docs/stable/generated/torch.floor_divide.html
+    // PyTorch aten.floor_divide is a misnomer because it actually rounds
+    // the quotient towards zero instead of taking its floor.
     Value cstStrFloor =
-        rewriter.create<Torch::ConstantStrOp>(op.getLoc(), "floor");
+        rewriter.create<Torch::ConstantStrOp>(op.getLoc(), "trunc");
     rewriter.replaceOpWithNewOp<AtenDivTensorModeOp>(
         op, op.getType(), op.self(), op.other(),
         /*rounding_mode=*/cstStrFloor);

test/Dialect/Torch/decompose-complex-ops.mlir

@@ -1113,8 +1113,8 @@ func.func @torch.aten.baddbmm(%arg0: !torch.vtensor<[?,?,?],f32>, %arg1: !torch.
 // CHECK-LABEL:   func @torch.aten.floor_divide(
 // CHECK-SAME:                  %[[SELF:.*]]: !torch.vtensor<[?,?],f32>,
 // CHECK-SAME:                  %[[OTHER:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
-// CHECK:           %[[CSTFLOOR:.*]] = torch.constant.str "floor"
-// CHECK:           %[[OUT:.*]] = torch.aten.div.Tensor_mode %[[SELF]], %[[OTHER]], %[[CSTFLOOR]] : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32>, !torch.str -> !torch.vtensor<[?,?],f32>
+// CHECK:           %[[CSTTRUNC:.*]] = torch.constant.str "trunc"
+// CHECK:           %[[OUT:.*]] = torch.aten.div.Tensor_mode %[[SELF]], %[[OTHER]], %[[CSTTRUNC]] : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32>, !torch.str -> !torch.vtensor<[?,?],f32>
 // CHECK:           return %[[OUT]] : !torch.vtensor<[?,?],f32>
 func.func @torch.aten.floor_divide(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
   %0 = torch.aten.floor_divide %arg0, %arg1 : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>