diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td index 02e7ed2c7..3e92d4099 100644 --- a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td +++ b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td @@ -14275,6 +14275,31 @@ def Torch_Aten_SetItemTOp : Torch_Op<"aten._set_item.t", [ }]; } +def Torch_AtenMulOp : Torch_Op<"aten.mul", [ + AllowsTypeRefinement, + HasValueSemantics, + ReadOnly + ]> { + let summary = "Generated op for `aten::mul : (Scalar, Scalar) -> (Scalar)`"; + let arguments = (ins + AnyTorchScalarType:$a, + AnyTorchScalarType:$b + ); + let results = (outs + AnyTorchScalarType:$result + ); + let hasCustomAssemblyFormat = 1; + let extraClassDefinition = [{ + ParseResult AtenMulOp::parse(OpAsmParser &parser, OperationState &result) { + return parseDefaultTorchOp(parser, result, 2, 1); + } + void AtenMulOp::print(OpAsmPrinter &printer) { + printDefaultTorchOp(printer, *this, 2, 1); + } + }]; + let hasFolder = 1; +} + def Torch_AtenDivOp : Torch_Op<"aten.div", [ AllowsTypeRefinement, HasValueSemantics, diff --git a/lib/Dialect/Torch/IR/TorchOps.cpp b/lib/Dialect/Torch/IR/TorchOps.cpp index 9f10c8bce..db2988f25 100644 --- a/lib/Dialect/Torch/IR/TorchOps.cpp +++ b/lib/Dialect/Torch/IR/TorchOps.cpp @@ -3346,6 +3346,25 @@ OpFoldResult AtenAddOp::fold(FoldAdaptor adaptor) { [](double a, double b) -> double { return a + b; }); } +//===----------------------------------------------------------------------===// +// AtenMulOp +//===----------------------------------------------------------------------===// + +OpFoldResult AtenMulOp::fold(FoldAdaptor adaptor) { + if (!adaptor.getA() || !adaptor.getB()) { + return nullptr; + } + + if (adaptor.getA().isa() && adaptor.getB().isa()) { + return atenBinaryIntOperatorFoldHelper( + adaptor.getOperands(), + [](int64_t a, int64_t b) -> int64_t { return a * b; }); + } + return atenBinaryFloatOperatorFoldHelper( + adaptor.getOperands(), + [](double a, double b) -> double { return a * b; }); +} + //===----------------------------------------------------------------------===// // AtenSubOp //===----------------------------------------------------------------------===// 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 54cdbea6a..055f7127c 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 @@ -805,6 +805,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry): has_canonicalizer=True) emit("aten::__getitem__.t : (t[], int) -> (t)", has_canonicalizer=True) emit("aten::_set_item.t : (t[], int, t) -> (t[])") + emit("aten::mul : (Scalar, Scalar) -> (Scalar)", has_folder=True) emit("aten::div : (Scalar, Scalar) -> (float)", has_folder=True) emit("aten::add : (Scalar, Scalar) -> (Scalar)", has_folder=True) emit("aten::sub : (Scalar, Scalar) -> (Scalar)", has_folder=True)