From d4862ec611ae728e3ed41ccbcab86e9946bd21c7 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 9 Dec 2022 20:52:26 +0530 Subject: [PATCH] [MLIR][TORCH] Add e2e support for aten.var_mean op Signed-Off By: Vivek Khandelwal --- e2e_testing/xfail_sets.py | 2 + .../Dialect/Torch/IR/GeneratedTorchOps.td | 25 +++++++++++ .../Torch/Transforms/DecomposeComplexOps.cpp | 21 +++++++++ .../Transforms/LowerToBackendContract.cpp | 1 + lib/Dialect/Torch/Transforms/RefineTypes.cpp | 2 +- lib/Dialect/Torch/Transforms/ShapeLibrary.cpp | 6 +++ .../jit_ir/build_tools/shape_lib_gen.py | 7 ++- .../jit_ir/build_tools/torch_ods_gen.py | 1 + .../torch_mlir_e2e_test/test_suite/stats.py | 44 +++++++++++++++++++ 9 files changed, 106 insertions(+), 3 deletions(-) diff --git a/e2e_testing/xfail_sets.py b/e2e_testing/xfail_sets.py index 80ad37f6d..a61e2d09e 100644 --- a/e2e_testing/xfail_sets.py +++ b/e2e_testing/xfail_sets.py @@ -811,4 +811,6 @@ LTC_XFAIL_SET = { "NllLossModule_sum_basic", "ResNet18Module_basic", "ResNet18StaticModule_basic", + "VarMeanBiasedModule_basic", + "VarMeanUnbiasedModule_basic", } diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td index 1c9651a9a..42d763a61 100644 --- a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td +++ b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td @@ -5008,6 +5008,31 @@ def Torch_AtenVarMeanCorrectionOp : Torch_Op<"aten.var_mean.correction", [ }]; } +def Torch_AtenVarMeanOp : Torch_Op<"aten.var_mean", [ + AllowsTypeRefinement, + HasValueSemantics, + ReadOnly + ]> { + let summary = "Generated op for `aten::var_mean : (Tensor, bool) -> (Tensor, Tensor)`"; + let arguments = (ins + AnyTorchTensorType:$self, + Torch_BoolType:$unbiased + ); + let results = (outs + AnyTorchTensorType:$result0, + AnyTorchTensorType:$result1 + ); + let hasCustomAssemblyFormat = 1; + let extraClassDefinition = [{ + ParseResult AtenVarMeanOp::parse(OpAsmParser &parser, OperationState &result) { + return parseDefaultTorchOp(parser, result, 2, 2); + } + void AtenVarMeanOp::print(OpAsmPrinter &printer) { + printDefaultTorchOp(printer, *this, 2, 2); + } + }]; +} + def Torch_AtenNllLossForwardOp : Torch_Op<"aten.nll_loss_forward", [ AllowsTypeRefinement, HasValueSemantics, diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp index 7107781b3..d0c3243f3 100644 --- a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp +++ b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp @@ -3284,6 +3284,26 @@ public: }; } // namespace +namespace { +class DecomposeAtenVarMeanOp : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(AtenVarMeanOp op, + PatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + Value falseVal = rewriter.create(loc, false); + Value noneVal = rewriter.create(loc); + Value var = rewriter.create(loc, op.getType(0), op.getSelf(), + /*dim=*/noneVal, op.getUnbiased(), + /*keepdim=*/falseVal); + Value mean = rewriter.create(loc, op.getType(0), op.getSelf(), + /*dtype=*/noneVal); + rewriter.replaceOp(op, {var, mean}); + return success(); + } +}; +} // namespace + namespace { class DecomposeComplexOpsPass : public DecomposeComplexOpsBase { @@ -3428,6 +3448,7 @@ public: addPatternIfTargetOpIsIllegal(patterns); addPatternIfTargetOpIsIllegal(patterns); addPatternIfTargetOpIsIllegal(patterns); + addPatternIfTargetOpIsIllegal(patterns); GreedyRewriteConfig config; config.useTopDownTraversal = true; diff --git a/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp b/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp index fa853796a..94493aa4e 100644 --- a/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp +++ b/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp @@ -418,6 +418,7 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context, target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); + target.addIllegalOp(); for (std::string opName : backendLegalOps) { target.addLegalOp(OperationName(opName, context)); } diff --git a/lib/Dialect/Torch/Transforms/RefineTypes.cpp b/lib/Dialect/Torch/Transforms/RefineTypes.cpp index 92db3cc44..770d92a35 100644 --- a/lib/Dialect/Torch/Transforms/RefineTypes.cpp +++ b/lib/Dialect/Torch/Transforms/RefineTypes.cpp @@ -1182,7 +1182,7 @@ void TypeAnalysis::visitOperation(Operation *op, return; } - if (isa(op)) { + if (isa(op)) { auto input = operands[0]->getValue(); auto knowledge = ValueKnowledge::getTensorPessimisticValueState(op->getContext()); diff --git a/lib/Dialect/Torch/Transforms/ShapeLibrary.cpp b/lib/Dialect/Torch/Transforms/ShapeLibrary.cpp index 1c16fe6d2..6ad95f4ef 100644 --- a/lib/Dialect/Torch/Transforms/ShapeLibrary.cpp +++ b/lib/Dialect/Torch/Transforms/ShapeLibrary.cpp @@ -5780,6 +5780,12 @@ StringRef mlir::torch::Torch::getShapeLibrary() { " %2 = torch.prim.TupleConstruct %1, %1 : !torch.list, !torch.list -> !torch.tuple, list>\n" " return %2 : !torch.tuple, list>\n" " }\n" +" func.func @\"__torch_mlir_shape_fn.aten.var_mean\"(%arg0: !torch.list, %arg1: !torch.bool) -> !torch.tuple, list> {\n" +" %0 = torch.prim.ListConstruct : () -> !torch.list\n" +" %1 = torch.prim.ListConstruct : () -> !torch.list\n" +" %2 = torch.prim.TupleConstruct %0, %1 : !torch.list, !torch.list -> !torch.tuple, list>\n" +" return %2 : !torch.tuple, list>\n" +" }\n" " func.func @\"__torch_mlir_shape_fn.aten.std\"(%arg0: !torch.list, %arg1: !torch.bool) -> !torch.list {\n" " %0 = torch.prim.ListConstruct : () -> !torch.list\n" " return %0 : !torch.list\n" diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py index 606023029..15e338c1d 100644 --- a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py +++ b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py @@ -544,13 +544,16 @@ def aten〇var(self: List[int], unbiased: bool = True) -> List[int]: def aten〇var〇dim(self: List[int], dim: Optional[List[int]], unbiased: bool = True, keepdim: bool = False) -> List[int]: return upstream_shape_functions.sum_mean_dim(self, dim, keepdim, None) -def aten〇var〇correction(self: List[int], dim: Optional[List[int]], correction: Optional[int], keepdim: bool = False) -> List[int]: +def aten〇var〇correction(self: List[int], dim: Optional[List[int]] = None, correction: Optional[int] = None, keepdim: bool = False) -> List[int]: return upstream_shape_functions.sum_mean_dim(self, dim, keepdim, None) -def aten〇var_mean〇correction(self: List[int], dim: Optional[List[int]], correction: Optional[int], keepdim: bool = False) -> Tuple[List[int], List[int]]: +def aten〇var_mean〇correction(self: List[int], dim: Optional[List[int]] = None, correction: Optional[int] = None, keepdim: bool = False) -> Tuple[List[int], List[int]]: out = upstream_shape_functions.sum_mean_dim(self, dim, keepdim, None) return out, out +def aten〇var_mean(self: List[int], unbiased: bool = True) -> Tuple[List[int], List[int]]: + return [], [] + def aten〇std(self: List[int], unbiased: bool = True) -> List[int]: return [] diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py index 8048a1499..bfe89e11d 100644 --- a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py +++ b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py @@ -409,6 +409,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry): emit("aten::var.dim : (Tensor, int[]?, bool, bool) -> (Tensor)") emit("aten::var.correction : (Tensor, int[]?, int?, bool) -> (Tensor)") emit("aten::var_mean.correction : (Tensor, int[]?, int?, bool) -> (Tensor, Tensor)") + emit("aten::var_mean : (Tensor, bool) -> (Tensor, Tensor)") emit("aten::nll_loss_forward : (Tensor, Tensor, Tensor?, int, int) -> (Tensor, Tensor)") emit("aten::nll_loss_backward : (Tensor, Tensor, Tensor, Tensor?, int, int, Tensor) -> (Tensor)") emit("aten::bincount : (Tensor, Tensor?, int) -> (Tensor)") diff --git a/python/torch_mlir_e2e_test/test_suite/stats.py b/python/torch_mlir_e2e_test/test_suite/stats.py index 3646b30c0..28112941f 100644 --- a/python/torch_mlir_e2e_test/test_suite/stats.py +++ b/python/torch_mlir_e2e_test/test_suite/stats.py @@ -799,3 +799,47 @@ class VarMeanCorrectionNoneModule(torch.nn.Module): @register_test_case(module_factory=lambda: VarMeanCorrectionNoneModule()) def VarMeanCorrectionNoneModule_basic(module, tu: TestUtils): module.forward(tu.rand(3, 4, 7)) + + +# ============================================================================== + + +class VarMeanUnbiasedModule(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.var_mean(x) + + +@register_test_case(module_factory=lambda: VarMeanUnbiasedModule()) +def VarMeanUnbiasedModule_basic(module, tu: TestUtils): + module.forward(tu.rand(3, 4, 7)) + + +# ============================================================================== + + +class VarMeanBiasedModule(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.var_mean(x, unbiased=False) + + +@register_test_case(module_factory=lambda: VarMeanBiasedModule()) +def VarMeanBiasedModule_basic(module, tu: TestUtils): + module.forward(tu.rand(3, 4, 7))