added the code for logcumsumexpOp

Author: sharavana20

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

@@ -8625,6 +8625,54 @@ def Torch_AtenCumprodOp : Torch_Op<"aten.cumprod", [
   }];
 }
 
+def Torch_AtenLogcumsumexpOp : Torch_Op<"aten.logcumsumexp", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::logcumsumexp : (Tensor, int) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$dim
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenLogcumsumexpOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenLogcumsumexpOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+}
+
+def Torch_Aten_LogcumsumexpOp : Torch_Op<"aten._logcumsumexp", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::_logcumsumexp : (Tensor, int) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$dim
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult Aten_LogcumsumexpOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void Aten_LogcumsumexpOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+}
+
 def Torch_AtenFloorDivideScalarOp : Torch_Op<"aten.floor_divide.Scalar", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -9361,6 +9361,9 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "  func.func @\"__torch_mlir_shape_fn.aten.cumprod\"(%arg0: !torch.list<int>, %arg1: !torch.int, %arg2: !torch.optional<int>) -> !torch.list<int> {\n"
 "    return %arg0 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.logcumsumexp\"(%arg0: !torch.list<int>, %arg1: !torch.int) -> !torch.list<int> {\n"
+"    return %arg0 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.rand_like\"(%arg0: !torch.list<int>, %arg1: !torch.optional<int>, %arg2: !torch.optional<int>, %arg3: !torch.optional<Device>, %arg4: !torch.optional<bool>, %arg5: !torch.optional<int>) -> !torch.list<int> {\n"
 "    return %arg0 : !torch.list<int>\n"
 "  }\n"
@@ -12507,6 +12510,11 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    }\n"
 "    return %1 : !torch.int\n"
 "  }\n"
+"  func.func @\"__torch_mlir_dtype_fn.aten.logcumsumexp\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int) -> !torch.int {\n"
+"    %int1 = torch.constant.int 1\n"
+"    %0 = torch.prim.TupleIndex %arg0, %int1 : !torch.tuple<int, int>, !torch.int -> !torch.int\n"
+"    return %0 : !torch.int\n"
+"  }\n"
 "  func.func @\"__torch_mlir_dtype_fn.aten.cumprod\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int, %arg2: !torch.optional<int>) -> !torch.int {\n"
 "    %int4 = torch.constant.int 4\n"
 "    %none = torch.constant.none\n"

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -2884,6 +2884,65 @@ class DecomposeAten_LogSoftmaxOp : public OpRewritePattern<Aten_LogSoftmaxOp> {
 };
 } // namespace
 
+// Decompose AtenLogCumsumExpOp to:
+// AtenExpOp
+// AtenCumsumOp
+// AtenLogOp
+namespace {
+
+class DecomposeAtenLogCumsumExpOp
+    : public OpRewritePattern<AtenLogcumsumexpOp> {
+public:
+  using OpRewritePattern<AtenLogcumsumexpOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenLogcumsumexpOp op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    Value input = op.getSelf();
+
+    auto inputType = dyn_cast<BaseTensorType>(input.getType());
+    if (!inputType || !inputType.getDtype())
+      return rewriter.notifyMatchFailure(op, "Supports only tensor type");
+
+    if (!inputType.hasDtype() || !isa<mlir::FloatType>(inputType.getDtype()))
+      return rewriter.notifyMatchFailure(op, "Support only floating type");
+
+    Type elementType = inputType.getDtype();
+    torch_upstream::ScalarType scalarType;
+    // logcumsumexp is only supported for Float datatype
+    if (elementType.isF16())
+      scalarType = torch_upstream::ScalarType::Half;
+    else if (elementType.isF32())
+      scalarType = torch_upstream::ScalarType::Float;
+    else
+      scalarType = torch_upstream::ScalarType::Double;
+
+    int64_t scalarVal = static_cast<int64_t>(scalarType);
+
+    Value dtypeVal = rewriter.create<Torch::ConstantIntOp>(
+        loc, rewriter.getType<Torch::IntType>(), scalarVal);
+
+    int64_t inputRank = inputType.getSizes().size();
+    int64_t dim;
+    if (!matchPattern(op.getDim(), m_TorchConstantInt(&dim)))
+      return rewriter.notifyMatchFailure(
+          op, "Only constant dim value is supported");
+    dim = toPositiveDim(dim, inputRank);
+    if (!isValidDim(dim, inputRank))
+      return rewriter.notifyMatchFailure(op, "invalid dim");
+
+    Value expInput = rewriter.create<AtenExpOp>(loc, input.getType(), input);
+
+    Value cumsum = rewriter.create<AtenCumsumOp>(
+        loc, expInput.getType(), expInput, op.getDim(), dtypeVal);
+
+    Value result = rewriter.create<AtenLogOp>(loc, cumsum.getType(), cumsum);
+
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class DecomposeAtenLogSigmoidOp : public OpRewritePattern<AtenLogSigmoidOp> {
 public:
@@ -11929,6 +11988,7 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAten_LogSoftmaxOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenLogSoftmaxIntOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenLogSigmoidOp>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenLogCumsumExpOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenHardshrinkOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenSoftshrinkOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenEmptyLikeOp>(patterns);

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -375,6 +375,7 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
   target.addIllegalOp<Aten_LogSoftmaxOp>();
   target.addIllegalOp<AtenLogSoftmaxIntOp>();
   target.addIllegalOp<AtenLogSigmoidOp>();
+  target.addIllegalOp<AtenLogcumsumexpOp>();
   target.addIllegalOp<AtenHardshrinkOp>();
   target.addIllegalOp<AtenSoftshrinkOp>();
   target.addIllegalOp<AtenEmptyLikeOp>();

projects/pt1/e2e_testing/xfail_sets.py

@@ -2965,6 +2965,10 @@
     "LinalgNormKeepDimComplexModule_basic",
     "LinalgVectorNormComplexModule_basic",
     "LogSoftmaxBackwardModule_basic",
+    "LogCumsumExpModule_basic",
+    "LogCumsumExpStaticModule_basic",
+    "LogCumsumExpStaticNegativeDimModule_basic",
+    "LogCumsumExpDtypeModule_basic",
     "MaxPool1dCeilModeTrueModule_basic",
     "MaxPool1dModule_basic",
     "MaxPool2dCeilModeTrueModule_basic",
@@ -3683,6 +3687,10 @@
     "LinalgNormKeepDimComplexModule_basic",
     "LinalgVectorNormComplexModule_basic",
     "LinspaceEmptyModule_basic",
+    "LogCumsumExpModule_basic",
+    "LogCumsumExpStaticModule_basic",
+    "LogCumsumExpStaticNegativeDimModule_basic",
+    "LogCumsumExpDtypeModule_basic",
     "MaskedScatterStaticBasic_basic",
     "MaxPool1dCeilModeTrueModule_basic",
     "MaxPool1dModule_basic",
@@ -4472,6 +4480,10 @@
     "LinalgVectorNormComplexModule_basic",
     "LogSoftmaxBackwardModule_basic",
     "LogSoftmaxIntModule_basic",
+    "logCumsumExpModule_basic",
+    "LogCumsumExpStaticModule_basic",
+    "LogCumsumExpStaticNegativeDimModule_basic",
+    "LogCumsumExpDtypeModule_basic",
     "MaskedFillTensorFloatValueModule_basic",
     "MatmulBroadcastBatchDim_basic",
     "MatmulSingleDynamicBatchDim_basic",

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

@@ -1537,6 +1537,9 @@ def aten〇cumsum〡shape(self: List[int], dim: int, dtype: Optional[int] = None
 def aten〇cumprod〡shape(self: List[int], dim: int, dtype: Optional[int] = None) -> List[int]:
     return self
 
+def aten〇logcumsumexp〡shape(self: List[int], dim: int) -> List[int]:
+    return self
+
 def aten〇rand_like〡shape(self: List[int], dtype: Optional[int] = None, layout: Optional[int] = None, device: Optional[device] = None, pin_memory: Optional[bool] = None, memory_format: Optional[int] = None) -> List[int]:
     return self
 
@@ -3217,6 +3220,10 @@ def aten〇cumsum〡dtype(self_rank_dtype: Tuple[int, int], dim: int, dtype: Opt
         return torch.int64
     return self_dtype
 
+@check_dtype_function(
+    _check_tensors_with_the_same_dtype(num_of_tensors=1, dim=0))
+def aten〇logcumsumexp〡dtype(self_rank_dtype: Tuple[int, int], dim: int) -> int:
+    return self_rank_dtype[1]
 
 @check_dtype_function(
     _check_tensors_with_the_same_dtype(num_of_tensors=1, dim=0) +

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -717,6 +717,8 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::bmm : (Tensor, Tensor) -> (Tensor)")
     emit("aten::cumsum : (Tensor, int, int?) -> (Tensor)")
     emit("aten::cumprod : (Tensor, int, int?) -> (Tensor)")
+    emit("aten::logcumsumexp : (Tensor, int) -> (Tensor)")
+    emit("aten::_logcumsumexp : (Tensor, int) -> (Tensor)")
     emit("aten::floor_divide.Scalar : (Tensor, Scalar) -> (Tensor)")
     emit("aten::logsumexp : (Tensor, int[], bool) -> (Tensor)")
     emit("aten::mean.dim : (Tensor, int[]?, bool, int?) -> (Tensor)")

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

@@ -5048,6 +5048,69 @@ def CumsumWithDtypeModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class LogCumsumExpModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([-1, -1, -1], torch.float32, True)])
+    def forward(self, x):
+        return torch.ops.aten.logcumsumexp(x, dim=1)
+
+
+@register_test_case(module_factory=lambda: LogCumsumExpModule())
+def LogCumsumExpModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(1, 2, 3))
+
+
+class LogCumsumExpStaticModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([1, 2, 3], torch.float32, True)])
+    def forward(self, x):
+        return torch.ops.aten.logcumsumexp(x, dim=1)
+
+
+@register_test_case(module_factory=lambda: LogCumsumExpStaticModule())
+def LogCumsumExpStaticModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(1, 2, 3))
+
+
+class LogCumsumExpStaticNegativeDimModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([8, 5, 6], torch.float32, True)])
+    def forward(self, x):
+        return torch.ops.aten.logcumsumexp(x, dim=-2)
+
+
+@register_test_case(module_factory=lambda: LogCumsumExpStaticNegativeDimModule())
+def LogCumsumExpStaticNegativeDimModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(8, 5, 6))
+
+
+class LogCumsumExpDtypeModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([5, 3, 6, 9], torch.float64, True)])
+    def forward(self, x):
+        return torch.ops.aten.logcumsumexp(x, dim=1)
+
+
+@register_test_case(module_factory=lambda: LogCumsumExpDtypeModule())
+def LogCumsumExpDtypeModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(5, 3, 6, 9).to(torch.float64))
+
+
+# ==============================================================================
+
+
 class CumprodModule(torch.nn.Module):
     def __init__(self):
         super().__init__()