[MLIR][TORCH] Add dtype arg support for torch.cumsum op (#4155)

This commit adds the support for dtype argument of torch.cumsum op which
specifies the desired data type of returned tensor.

Fixes #3866.

---------

Signed-off-by: Vivek Khandelwal <[email protected]>

Author: vivekkhandelwal1

lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp

@@ -1634,14 +1634,35 @@ class ConvertAtenCumsumOp : public OpConversionPattern<AtenCumsumOp> {
     auto resultType = cast<RankedTensorType>(
         getTypeConverter()->convertType(op->getResult(0).getType()));
     Type elementType = resultType.getElementType();
-    Type inputElementType =
-        cast<RankedTensorType>(input.getType()).getElementType();
+    auto inputType = cast<RankedTensorType>(input.getType());
+    Type inputElementType = inputType.getElementType();
 
-    // Converting the input element type to the result's element type.
-    // The only possible mismatch would be when the input element type is an
-    // integer but not `si64`. Therefore, we directly convert the input to
-    // `si64`. Rest all cases are handled in the dtype definition for this op.
-    if (elementType != inputElementType) {
+    Value dtype = op.getDtype();
+    if (!isa<Torch::NoneType>(dtype.getType())) {
+      int64_t dtypeInt;
+      if (!matchPattern(dtype, m_TorchConstantInt(&dtypeInt)))
+        return rewriter.notifyMatchFailure(
+            op, "unimplemented: only constant int dtype value is supported");
+
+      FailureOr<Type> resDtype = getTypeForScalarType(
+          op->getContext(), (torch_upstream::ScalarType)dtypeInt);
+      if (failed(resDtype))
+        return rewriter.notifyMatchFailure(
+            op, "unsupported: dtype not defined for the given dtype int value");
+
+      Value torchInput =
+          convertTensorToDtype(rewriter, loc, op.getSelf(), resDtype.value());
+      input = typeConverter->materializeTargetConversion(
+          rewriter, loc, typeConverter->convertType(torchInput.getType()),
+          torchInput);
+    } else if (elementType != inputElementType &&
+               isa<mlir::IntegerType>(elementType) &&
+               isa<mlir::IntegerType>(inputElementType)) {
+      // Converting the input element type to the result's element type.
+      // The only possible mismatch would be when the input element type is an
+      // integer but not `si64` and the `dtype` is not specified. Therefore, we
+      // directly convert the input to `si64`. Rest all cases are handled in the
+      // dtype definition for this op.
       Value torchInput = convertTensorToDtype(
           rewriter, loc, op.getSelf(),
           rewriter.getIntegerType(64, IntegerType::Signed));
@@ -1650,12 +1671,7 @@ class ConvertAtenCumsumOp : public OpConversionPattern<AtenCumsumOp> {
           torchInput);
     }
 
-    int64_t inputRank = resultType.getRank();
-    Value dtype = op.getDtype();
-    if (!isa<Torch::NoneType>(dtype.getType()))
-      return rewriter.notifyMatchFailure(
-          op, "unsupported: dtype argument not supported");
-
+    int64_t inputRank = inputType.getRank();
     int64_t dim;
     if (!matchPattern(op.getDim(), m_TorchConstantInt(&dim)))
       return rewriter.notifyMatchFailure(

projects/pt1/e2e_testing/xfail_sets.py

@@ -3579,6 +3579,7 @@
     "ConvolutionModule2DTranspose_basic",
     "ConvolutionModule2DGroupedTranspose_basic",
     "CumsumInputDtypeInt32Module_basic",
+    "CumsumWithDtypeModule_basic",
     "CumsumModule_basic",
     "CumsumStaticModule_basic",
     "CumsumStaticNegativeDimModule_basic",

projects/pt1/python/torch_mlir_e2e_test/test_suite/basic.py

@@ -5025,6 +5025,26 @@ def CumsumInputDtypeInt32Module_basic(module, tu: TestUtils):
     module.forward(tu.randint(2, 7, 4).to(torch.int32))
 
 
+class CumsumWithDtypeModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([2, 7, 4], torch.bool, True),
+        ]
+    )
+    def forward(self, val):
+        return torch.ops.aten.cumsum(val, dim=1, dtype=6)
+
+
+@register_test_case(module_factory=lambda: CumsumWithDtypeModule())
+def CumsumWithDtypeModule_basic(module, tu: TestUtils):
+    module.forward(tu.randint(2, 7, 4, low=-1, high=10).to(torch.bool))
+
+
 # ==============================================================================