From abef114c0c067b78b973cae17e687add5138e969 Mon Sep 17 00:00:00 2001 From: Xinyu Yang Date: Wed, 8 May 2024 15:20:45 +0800 Subject: [PATCH] [torch] emit aten.Softshrink and aten.Hardshrink (#3248) as title --- .../Dialect/Torch/IR/GeneratedTorchOps.td | 48 ++++++++ .../Transforms/AbstractInterpLibrary.cpp | 25 ++++ .../Torch/Transforms/DecomposeComplexOps.cpp | 116 ++++++++++++++++++ .../Transforms/LowerToBackendContract.cpp | 2 + projects/pt1/e2e_testing/xfail_sets.py | 6 + .../build_tools/abstract_interp_lib_gen.py | 18 +++ .../build_tools/torch_ods_gen.py | 2 + .../test_suite/elementwise.py | 92 ++++++++++++++ 8 files changed, 309 insertions(+) diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td index 95d92af992b9..662fb8f0d5bd 100644 --- a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td +++ b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td @@ -5019,6 +5019,54 @@ def Torch_AtenLogSigmoidOp : Torch_Op<"aten.log_sigmoid", [ }]; } +def Torch_AtenHardshrinkOp : Torch_Op<"aten.hardshrink", [ + AllowsTypeRefinement, + HasValueSemantics, + ReadOnly + ]> { + let summary = "Generated op for `aten::hardshrink : (Tensor, Scalar) -> (Tensor)`"; + let arguments = (ins + AnyTorchTensorType:$self, + AnyTorchScalarType:$lambd + ); + let results = (outs + AnyTorchOptionalTensorType:$result + ); + let hasCustomAssemblyFormat = 1; + let extraClassDefinition = [{ + ParseResult AtenHardshrinkOp::parse(OpAsmParser &parser, OperationState &result) { + return parseDefaultTorchOp(parser, result, 2, 1); + } + void AtenHardshrinkOp::print(OpAsmPrinter &printer) { + printDefaultTorchOp(printer, *this, 2, 1); + } + }]; +} + +def Torch_AtenSoftshrinkOp : Torch_Op<"aten.softshrink", [ + AllowsTypeRefinement, + HasValueSemantics, + ReadOnly + ]> { + let summary = "Generated op for `aten::softshrink : (Tensor, Scalar) -> (Tensor)`"; + let arguments = (ins + AnyTorchTensorType:$self, + AnyTorchScalarType:$lambd + ); + let results = (outs + AnyTorchOptionalTensorType:$result + ); + let hasCustomAssemblyFormat = 1; + let extraClassDefinition = [{ + ParseResult AtenSoftshrinkOp::parse(OpAsmParser &parser, OperationState &result) { + return parseDefaultTorchOp(parser, result, 2, 1); + } + void AtenSoftshrinkOp::print(OpAsmPrinter &printer) { + printDefaultTorchOp(printer, *this, 2, 1); + } + }]; +} + def Torch_AtenUnbindCopyIntOp : Torch_Op<"aten.unbind_copy.int", [ AllowsTypeRefinement, HasValueSemantics, diff --git a/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp b/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp index d9ac7a6d0c55..259df7ef6890 100644 --- a/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp +++ b/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp @@ -6514,6 +6514,14 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() { " %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list\n" " return %0 : !torch.list\n" " }\n" +" func.func @\"__torch_mlir_shape_fn.aten.hardshrink\"(%arg0: !torch.list, %arg1: !torch.float) -> !torch.list {\n" +" %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list\n" +" return %0 : !torch.list\n" +" }\n" +" func.func @\"__torch_mlir_shape_fn.aten.softshrink\"(%arg0: !torch.list, %arg1: !torch.float) -> !torch.list {\n" +" %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list\n" +" return %0 : !torch.list\n" +" }\n" " func.func @\"__torch_mlir_shape_fn.aten.mish\"(%arg0: !torch.list) -> !torch.list {\n" " %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list\n" " return %0 : !torch.list\n" @@ -9798,6 +9806,23 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() { " }\n" " return %0#1 : !torch.int\n" " }\n" +" func.func @\"__torch_mlir_dtype_fn.aten.hardshrink\"(%arg0: !torch.tuple, %arg1: !torch.number) -> !torch.int {\n" +" %int4 = torch.constant.int 4\n" +" %int11 = torch.constant.int 11\n" +" %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple -> !torch.int, !torch.int\n" +" %1 = torch.aten.eq.int %0#1, %int11 : !torch.int, !torch.int -> !torch.bool\n" +" %2 = torch.prim.If %1 -> (!torch.int) {\n" +" torch.prim.If.yield %int4 : !torch.int\n" +" } else {\n" +" torch.prim.If.yield %0#1 : !torch.int\n" +" }\n" +" return %2 : !torch.int\n" +" }\n" +" func.func @\"__torch_mlir_dtype_fn.aten.softshrink\"(%arg0: !torch.tuple, %arg1: !torch.number) -> !torch.int {\n" +" %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple -> !torch.int, !torch.int\n" +" %1 = call @__torch__._get_dtype_of_floating_point_op(%0#1) : (!torch.int) -> !torch.int\n" +" return %1 : !torch.int\n" +" }\n" " func.func @\"__torch_mlir_dtype_fn.aten.logit\"(%arg0: !torch.tuple, %arg1: !torch.optional) -> !torch.int {\n" " %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple -> !torch.int, !torch.int\n" " %1 = call @__torch__._get_dtype_of_floating_point_op(%0#1) : (!torch.int) -> !torch.int\n" diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp index cc21f2155e46..73a67e14b132 100644 --- a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp +++ b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp @@ -1913,6 +1913,120 @@ class DecomposeAtenLogSigmoidOp : public OpRewritePattern { }; } // namespace +// SoftShrink(x, lambda) function: +// Applies a shrinkage function where: +// - If x > lambda, returns x - lambda +// - If x < -lambda, returns x + lambda +// - Otherwise, returns 0 +namespace { +class DecomposeAtenSoftshrinkOp : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(AtenSoftshrinkOp op, + PatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + Value self = op.getSelf(); + Value lambdValue = op.getLambd(); + + auto resTy = dyn_cast(op.getType()); + if (!resTy || !resTy.hasDtype() || !resTy.hasSizes()) { + return rewriter.notifyMatchFailure(op, + "result should have dtype and size"); + } + + double lambd; + if (!matchPattern(lambdValue, m_TorchConstantFloat(&lambd))) { + return rewriter.notifyMatchFailure( + op, "expected lambd to be a constant float"); + } + + Value zero = + rewriter.create(loc, rewriter.getF64FloatAttr(0.0)); + Value neglambd = rewriter.create( + loc, rewriter.getF64FloatAttr(-lambd)); + Value poslambd = rewriter.create( + loc, rewriter.getF64FloatAttr(lambd)); + + Value constOneFloat = + rewriter.create(loc, rewriter.getF64FloatAttr(1.0)); + + auto boolResType = + resTy.getWithSizesAndDtype(resTy.getSizes(), rewriter.getI1Type()); + + Value posMask = + rewriter.create(loc, boolResType, self, poslambd); + Value negMask = + rewriter.create(loc, boolResType, self, neglambd); + + Value posValue = rewriter.create(loc, resTy, self, + poslambd, constOneFloat); + Value negValue = rewriter.create(loc, resTy, self, + neglambd, constOneFloat); + + Value result = rewriter.create(loc, resTy, posMask, + posValue, zero); + result = + rewriter.create(loc, resTy, negMask, negValue, result); + + rewriter.replaceOp(op, result); + return success(); + } +}; +} // namespace + +// HardShrink(x, lambda) function: +// Applies a shrinkage function where: +// - If x > lambda, returns x +// - If x < -lambda, returns x +// - Otherwise, returns 0 +namespace { +class DecomposeAtenHardshrinkOp : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(AtenHardshrinkOp op, + PatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + Value self = op.getSelf(); + Value lambdValue = op.getLambd(); + + auto resTy = dyn_cast(op.getType()); + if (!resTy || !resTy.hasDtype() || !resTy.hasSizes()) { + return rewriter.notifyMatchFailure(op, + "result should have dtype and size"); + } + + double lambd; + if (!matchPattern(lambdValue, m_TorchConstantFloat(&lambd))) { + return rewriter.notifyMatchFailure( + op, "expected lambd to be a constant float"); + } + + Value zero = + rewriter.create(loc, rewriter.getF64FloatAttr(0.0)); + Value neglambd = rewriter.create( + loc, rewriter.getF64FloatAttr(-lambd)); + Value poslambd = rewriter.create( + loc, rewriter.getF64FloatAttr(lambd)); + + auto boolResType = + resTy.getWithSizesAndDtype(resTy.getSizes(), rewriter.getI1Type()); + + Value posMask = + rewriter.create(loc, boolResType, self, poslambd); + Value negMask = + rewriter.create(loc, boolResType, self, neglambd); + + Value result = rewriter.create(loc, resTy, posMask, + self, zero); + result = + rewriter.create(loc, resTy, negMask, self, result); + + rewriter.replaceOp(op, result); + return success(); + } +}; +} // namespace + // Decompose aten.matmul into: aten.mm and aten.bmm according to ranks. namespace { class DecomposeAtenMatmulOp : public OpRewritePattern { @@ -7664,6 +7778,8 @@ class DecomposeComplexOpsPass addPatternIfTargetOpIsIllegal(patterns); addPatternIfTargetOpIsIllegal(patterns); addPatternIfTargetOpIsIllegal(patterns); + addPatternIfTargetOpIsIllegal(patterns); + addPatternIfTargetOpIsIllegal(patterns); addPatternIfTargetOpIsIllegal(patterns); addPatternIfTargetOpIsIllegal< DecomposeConstantTensorAllocLikeOp>(patterns); diff --git a/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp b/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp index e7bed6463552..36dc4071161f 100644 --- a/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp +++ b/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp @@ -371,6 +371,8 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context, target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); + target.addIllegalOp(); + target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py index fca01708bceb..c3a300f36908 100644 --- a/projects/pt1/e2e_testing/xfail_sets.py +++ b/projects/pt1/e2e_testing/xfail_sets.py @@ -1438,6 +1438,8 @@ "ElementwiseTruncIntModule_basic", "ElementwiseTruncModule_basic", "ElementwiseLogSigmoidModule_basic", + "ElementwiseHardshrinkStaticModule_basic", + "ElementwiseSoftshrinkStaticModule_basic", } STABLEHLO_CRASHING_SET = { @@ -1667,6 +1669,10 @@ "ElementwiseSeluModule_basic", "ElementwiseSigmoidModule_basic", "ElementwiseSignModule_basic", + "ElementwiseHardshrinkModule_basic", + "ElementwiseHardshrinkStaticModule_basic", + "ElementwiseSoftshrinkModule_basic", + "ElementwiseSoftshrinkStaticModule_basic", "ElementwiseSqrtIntModule_basic", "ElementwiseSqrtModule_basic", "ElementwiseSubScalarFloatModule_basic", diff --git a/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py b/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py index eb60620561bf..7cb85f103c2c 100644 --- a/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py +++ b/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py @@ -254,6 +254,12 @@ def aten〇log〡shape(self: List[int]) -> List[int]: def aten〇log_sigmoid〡shape(self: List[int]) -> List[int]: return upstream_shape_functions.unary(self) +def aten〇hardshrink〡shape(self: List[int], lambd: float = 0.5) -> List[int]: + return upstream_shape_functions.unary(self) + +def aten〇softshrink〡shape(self: List[int], lambd: float = 0.5) -> List[int]: + return upstream_shape_functions.unary(self) + def aten〇mish〡shape(self: List[int]) -> List[int]: return upstream_shape_functions.unary(self) @@ -2098,6 +2104,18 @@ def aten〇log_sigmoid〡dtype(self_rank_dtype: Tuple[int, int]) -> int: assert not self_dtype == torch.bool return self_dtype +@check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1, lambd=0.5)) +def aten〇hardshrink〡dtype(self_rank_dtype: Tuple[int, int], lambd: Union[int, float, complex] = 0.5) -> int: + self_rank, self_dtype = self_rank_dtype + if self_dtype == torch.bool: + return torch.int64 + return self_dtype + +@check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1, lambd=0.5)) +def aten〇softshrink〡dtype(self_rank_dtype: Tuple[int, int], lambd: Union[int, float, complex] = 0.5) -> int: + self_rank, self_dtype = self_rank_dtype + return _get_dtype_of_floating_point_op(self_dtype) + @check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1)) def aten〇logit〡dtype(self_rank_dtype: Tuple[int, int], eps: Optional[float] = None) -> int: self_rank, self_dtype = self_rank_dtype diff --git a/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py b/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py index d4d547456c4a..515f829ff407 100644 --- a/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py +++ b/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py @@ -480,6 +480,8 @@ def emit_with_mutating_variants(key, **kwargs): emit("aten::isclose : (Tensor, Tensor, float, float, bool) -> (Tensor)") emit("aten::glu : (Tensor, int) -> (Tensor)") emit("aten::log_sigmoid : (Tensor) -> (Tensor)") + emit("aten::hardshrink : (Tensor, Scalar) -> (Tensor)") + emit("aten::softshrink : (Tensor, Scalar) -> (Tensor)") # Ops with dynamic number of outputs emit("aten::unbind_copy.int : (Tensor, int) -> (Tensor[])") diff --git a/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py b/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py index a26fd9809f13..f5e3c9fc4b9b 100644 --- a/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py +++ b/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py @@ -2205,6 +2205,98 @@ def ElementwiseLogSigmoidModule_basic(module, tu: TestUtils): # ============================================================================== +class ElementwiseSoftshrinkModule(torch.nn.Module): + def __init__(self): + super().__init__() + + @export + @annotate_args( + [ + None, + ([-1, -1], torch.float32, True), + ] + ) + def forward(self, a): + return torch.ops.aten.softshrink(a) + + +@register_test_case(module_factory=lambda: ElementwiseSoftshrinkModule()) +def ElementwiseSoftshrinkModule_basic(module, tu: TestUtils): + module.forward(tu.rand(3, 4)) + + +# ============================================================================== + + +class ElementwiseSoftshrinkStaticModule(torch.nn.Module): + def __init__(self): + super().__init__() + + @export + @annotate_args( + [ + None, + ([4, 5, 6], torch.float32, True), + ] + ) + def forward(self, a): + return torch.ops.aten.softshrink(a, 2.0) + + +@register_test_case(module_factory=lambda: ElementwiseSoftshrinkStaticModule()) +def ElementwiseSoftshrinkStaticModule_basic(module, tu: TestUtils): + module.forward(tu.rand(4, 5, 6)) + + +# ============================================================================== + + +class ElementwiseHardshrinkModule(torch.nn.Module): + def __init__(self): + super().__init__() + + @export + @annotate_args( + [ + None, + ([-1, -1, -1], torch.float32, True), + ] + ) + def forward(self, a): + return torch.ops.aten.hardshrink(a, 1.0) + + +@register_test_case(module_factory=lambda: ElementwiseHardshrinkModule()) +def ElementwiseHardshrinkModule_basic(module, tu: TestUtils): + module.forward(tu.rand(3, 4, 5)) + + +# ============================================================================== + + +class ElementwiseHardshrinkStaticModule(torch.nn.Module): + def __init__(self): + super().__init__() + + @export + @annotate_args( + [ + None, + ([4, 5, 6], torch.float32, True), + ] + ) + def forward(self, a): + return torch.ops.aten.hardshrink(a, 2.0) + + +@register_test_case(module_factory=lambda: ElementwiseHardshrinkStaticModule()) +def ElementwiseHardshrinkStaticModule_basic(module, tu: TestUtils): + module.forward(tu.rand(4, 5, 6)) + + +# ============================================================================== + + class ElementwiseErfModule(torch.nn.Module): def __init__(self): super().__init__()