[FXML-2172] Torch Div Conversion Fix (#45)

* fix(TorchToTosa.cpp): adjust torch div conversion

check the return type of the division to figure out whether to use
the floating point implementation of a division or to use the
integer.

the issue rose from the fact that the inputs are all integer but
the result was casted to floating point. The conversion then chose
to use the integer implementation of division which is not legal in
tosa when all the inputs get casted to floating point.

* test(e2e): integer division resulting in a float

pytorch example of two integers being divided that should case to a float

* fix(TorchToTosa.cpp): correct type promotion for reciprocal

the operation should only be handling floats and not integers

* Update python/torch_mlir_e2e_test/test_suite/elementwise.py

Co-authored-by: Matthias Gehre <93204396+mgehre-amd@users.noreply.github.com>

* fix(xfail_sets.py): add torchdynamo case

for tensor divided by scalar

---------

Co-authored-by: Matthias Gehre <93204396+mgehre-amd@users.noreply.github.com>

Author: cmcgirr-amd

e2e_testing/xfail_sets.py

@@ -229,6 +229,9 @@
     # ERROR: 'torch.aten.div.Tensor' op operand #1 must be Any Torch tensor type, but got '!torch.float'
     "ElementwiseDivScalarModule_basic",
 
+    # ERROR: 'torch.aten.div.Tensor' op operand #1 must be Any Torch tensor type, but got '!torch.int'
+    "ElementwiseDivIntScalarModule_basic",
+
     # ERROR: 'torch.aten.mul.Tensor' op operand #1 must be Any Torch tensor type, but got '!torch.int'
     "ElementwiseMulScalarModule_int",
 
@@ -424,6 +427,7 @@
     "ElementwiseAddScalar_NumToTensorFloat_Module_basic",
     "ElementwiseAddScalar_TensorLiteralInt32_Module_basic",
     "ElementwiseDivScalarModule_basic",
+    "ElementwiseDivIntScalarModule_basic",
     "ElementwiseEqDiffWidthScalarModule_basic",
     "ElementwiseEqFloatScalarModule_basic",
     "ElementwiseEqIntScalarModule_basic",
@@ -893,6 +897,7 @@
     "ElementwiseMulScalarModule_float",
     "ElementwiseMulTensorIntModule_basic",
     "ElementwiseDivScalarModule_basic",
+    "ElementwiseDivIntScalarModule_basic",
     "ElementwiseSubScalarFloatModule_basic",
     "ElementwiseAddScalarFloatModule_basic",
     "ElementwiseAddScalar_TensorLiteralInt32_Module_basic",

lib/Conversion/TorchToLinalg/Uncategorized.cpp

@@ -958,7 +958,7 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
       divScalar.emitError("unimplemented: non-floating point dtype");
       return nullptr;
     }
-    Value self = payloadArgs[0];
+    Value self = convertScalarToDtype(b, loc, payloadArgs[0], dtype);
     Value other = convertScalarToDtype(b, loc, operands[1], dtype);
     return b.create<arith::DivFOp>(loc, self, other);
   }

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -497,15 +497,20 @@ class ConvertAtenDivOp : public OpConversionPattern<AtenOpT> {
 
     // auto result;
     Value result;
-    if (lhsElemTy.isa<mlir::FloatType>()) {
+    if (outType.getElementType().template isa<mlir::FloatType>()) {
+      // The input to the reciprocal is an integer sometimes, and we may need to
+      // promote it to a floating point. Per TOSA specification, the input types
+      // can only be floating point for tosa::ReciprocalOp.
+      Value rhsCasted = tosa::promoteType(rewriter, rhsTensor, outType);
       auto rcpOp = rewriter.create<tosa::ReciprocalOp>(
-          op->getLoc(), rhsTy ? rhsTy : RankedTensorType::get({}, lhsElemTy),
-          rhsTensor);
+          op->getLoc(), rhsCasted.getType(), rhsCasted);
 
       result = tosa::createMulOpAndCast(rewriter, op, outType, lhs,
                                         rcpOp.getResult(), /*shift=*/0)
                    .getResult();
     } else {
+      // If the output type of the original operation is an integer then we will
+      // apply a tosa div knowing that rounding will occur and truncate to zero.
       result = tosa::createBinaryOpAndCast<tosa::DivOp>(rewriter, op, outType,
                                                         lhs, rhsTensor)
                    .getResult();

python/torch_mlir_e2e_test/test_suite/elementwise.py

@@ -1744,6 +1744,28 @@ def forward(self, x):
 def ElementwiseDivScalarModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
 
+
+# ==============================================================================
+
+
+class ElementwiseDivIntScalarModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.int64, True),
+    ])
+    def forward(self, x):
+        return torch.ops.aten.div(x, 128)
+
+
+@register_test_case(module_factory=lambda: ElementwiseDivIntScalarModule())
+def ElementwiseDivIntScalarModule_basic(module, tu: TestUtils):
+    module.forward(tu.randint(3, 4))
+
 # ==============================================================================
 
 

test/Conversion/TorchToTosa/torch-backend-to-tosa-backend-pipeline.mlir

@@ -91,8 +91,8 @@ func.func @torch.aten.bitwise_and.Tensor$mixed_type(%arg0: !torch.vtensor<[?,?],
 // CHECK-LABEL: torch.aten.div.Tensor$mixed_type_fp
 // CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xf32>,
 // CHECK-SAME: %[[VAL_1:.*]]: tensor<?x?xi32>
-// CHECK: %[[VAL_2:.*]] = "tosa.reciprocal"(%[[VAL_1]]) : (tensor<?x?xi32>) -> tensor<?x?xi32>
-// CHECK: %[[VAL_3:.*]] = "tosa.cast"(%[[VAL_2]]) : (tensor<?x?xi32>) -> tensor<?x?xf32>
+// CHECK: %[[VAL_2:.*]] = "tosa.cast"(%[[VAL_1]]) : (tensor<?x?xi32>) -> tensor<?x?xf32>
+// CHECK: %[[VAL_3:.*]] = "tosa.reciprocal"(%[[VAL_2]]) : (tensor<?x?xf32>) -> tensor<?x?xf32>
 // CHECK: %[[VAL_4:.*]] = "tosa.mul"(%[[VAL_0]], %[[VAL_3]]) <{shift = 0 : i32}> : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
 func.func @torch.aten.div.Tensor$mixed_type_fp(%arg0: !torch.vtensor<[?, ?],f32>, %arg1: !torch.vtensor<[?, ?],si32>) -> !torch.vtensor<[?, ?],f32> {
   %0 = torch.aten.div.Tensor %arg0, %arg1 : !torch.vtensor<[?, ?],f32>, !torch.vtensor<[?, ?],si32> -> !torch.vtensor<[?, ?],f32>
@@ -113,6 +113,21 @@ func.func @torch.aten.div.Tensor$mixed_type_int(%arg0: !torch.vtensor<[?, ?],si1
 
 // -----
 
+// CHECK-LABEL: torch.aten.div.Scalar$int_input_fp_output
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xi64>
+// CHECK: %[[VAL_1:.*]] = "tosa.const"() <{value = dense<1.280000e+02> : tensor<f32>}> : () -> tensor<f32>
+// CHECK: %[[VAL_2:.*]] = "tosa.reciprocal"(%[[VAL_1]]) : (tensor<f32>) -> tensor<f32>
+// CHECK: %[[VAL_3:.*]] = "tosa.cast"(%[[VAL_0]]) : (tensor<?x?xi64>) -> tensor<?x?xf32>
+// CHECK: %[[VAL_4:.*]] = "tosa.reshape"(%[[VAL_2]]) <{new_shape = array<i64: 1, 1>}> : (tensor<f32>) -> tensor<1x1xf32>
+// CHECK: %[[VAL_5:.*]] = "tosa.mul"(%[[VAL_3]], %[[VAL_4]]) <{shift = 0 : i32}> : (tensor<?x?xf32>, tensor<1x1xf32>) -> tensor<?x?xf32>
+func.func @torch.aten.div.Scalar$int_input_fp_output(%arg0: !torch.vtensor<[?, ?],si64>) -> !torch.vtensor<[?, ?],f32> {
+  %int128 = torch.constant.int 128
+  %0 = torch.aten.div.Scalar %arg0, %int128 : !torch.vtensor<[?, ?],si64>, !torch.int  -> !torch.vtensor<[?, ?],f32>
+  return %0 : !torch.vtensor<[?, ?],f32>
+}
+
+// -----
+
 // CHECK-LABEL: torch.aten.pow.Tensor$mixed_type
 // CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xf16>
 // CHECK: %[[VAL_1:.*]] = "tosa.const"() <{value = dense<3.000000e+00> : tensor<1x1xf32>}> : () -> tensor<1x1xf32>