[Torch Dialect] Emit torch.aten.mul.float_int, add folder and conversion to Arith.

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

@@ -15934,6 +15934,31 @@ def Torch_AtenAddFloatIntOp : Torch_Op<"aten.add.float_int", [
   let hasFolder = 1;
 }
 
+def Torch_AtenMulFloatIntOp : Torch_Op<"aten.mul.float_int", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::mul.float_int : (float, int) -> (float)`";
+  let arguments = (ins
+    Torch_FloatType:$a,
+    Torch_IntType:$b
+  );
+  let results = (outs
+    Torch_FloatType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenMulFloatIntOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenMulFloatIntOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+  let hasFolder = 1;
+}
+
 def Torch_AtenSubFloatOp : Torch_Op<"aten.sub.float", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Conversion/TorchToArith/TorchToArith.cpp

@@ -74,7 +74,8 @@ class ConvertAtenBinaryOp : public OpConversionPattern<AtenOp> {
                   ConversionPatternRewriter &rewriter) const override {
     Value a = adaptor.getA();
     Value b = adaptor.getB();
-    if (llvm::is_one_of<AtenOp, AtenAddFloatIntOp>::value)
+    if (llvm::is_one_of<AtenOp, AtenAddFloatIntOp>::value ||
+        llvm::is_one_of<AtenOp, AtenMulFloatIntOp>::value)
       b = convertScalarToDtype(rewriter, op.getLoc(), b, a.getType());
     rewriter.template replaceOpWithNewOp<BinOp>(op, a, b);
     return success();
@@ -487,7 +488,7 @@ class ConvertTorchToArith
     target.addIllegalOp<AtenNegIntOp>();
     patterns.add<ConvertAtenNegIntOp>(typeConverter, context);
     target.addIllegalOp<AtenAddIntOp, AtenAddFloatIntOp, AtenSubIntOp,
-                        AtenMulIntOp, AtenRemainderIntOp>();
+                        AtenMulIntOp, AtenRemainderIntOp, AtenMulFloatIntOp>();
     patterns.add<ConvertAtenBinaryOp<AtenAddIntOp, arith::AddIOp>>(
         typeConverter, context);
     patterns.add<ConvertAtenBinaryOp<AtenRemainderIntOp, arith::RemSIOp>>(
@@ -498,6 +499,8 @@ class ConvertTorchToArith
         typeConverter, context);
     patterns.add<ConvertAtenBinaryOp<AtenMulIntOp, arith::MulIOp>>(
         typeConverter, context);
+    patterns.add<ConvertAtenBinaryOp<AtenMulFloatIntOp, arith::MulFOp>>(
+        typeConverter, context);
     target.addIllegalOp<AtenSubFloatOp, AtenMulFloatOp>();
     patterns.add<ConvertAtenBinaryOp<AtenSubFloatOp, arith::SubFOp>>(
         typeConverter, context);

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -4220,6 +4220,18 @@ OpFoldResult AtenAddFloatIntOp::fold(FoldAdaptor adaptor) {
       adaptor.getOperands(), [](double a, double b) { return a + b; });
 }
 
+//===----------------------------------------------------------------------===//
+// AtenMulFloatIntOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AtenMulFloatIntOp::fold(FoldAdaptor adaptor) {
+  if (!adaptor.getA() || !adaptor.getB()) {
+    return nullptr;
+  }
+  return atenBinaryFloatOperatorFoldHelper(
+      adaptor.getOperands(), [](double a, double b) { return a * b; });
+}
+
 //===----------------------------------------------------------------------===//
 // AtenPowIntFloatOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -6908,12 +6908,12 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "        }\n"
 "        %11 = torch.aten.__getitem__.t %9, %int0 : !torch.list<float>, !torch.int -> !torch.float\n"
 "        %12 = torch.aten.__getitem__.t %arg0, %int2 : !torch.list<int>, !torch.int -> !torch.int\n"
-"        %13 = torch.operator \"aten.mul.float_int\"(%11, %12) : (!torch.float, !torch.int) -> !torch.float \n"
+"        %13 = torch.aten.mul.float_int %11, %12 : !torch.float, !torch.int -> !torch.float\n"
 "        %14 = torch.aten.Int.float %13 : !torch.float -> !torch.int\n"
 "        %15 = torch.aten.append.t %3, %14 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
 "        %16 = torch.aten.__getitem__.t %9, %int1 : !torch.list<float>, !torch.int -> !torch.float\n"
 "        %17 = torch.aten.__getitem__.t %arg0, %int3 : !torch.list<int>, !torch.int -> !torch.int\n"
-"        %18 = torch.operator \"aten.mul.float_int\"(%16, %17) : (!torch.float, !torch.int) -> !torch.float \n"
+"        %18 = torch.aten.mul.float_int %16, %17 : !torch.float, !torch.int -> !torch.float\n"
 "        %19 = torch.aten.Int.float %18 : !torch.float -> !torch.int\n"
 "        %20 = torch.aten.append.t %3, %19 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
 "        torch.prim.If.yield %true, %3 : !torch.bool, !torch.list<int>\n"

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

@@ -1120,6 +1120,7 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::neg.int : (int) -> (int)", has_folder=True)
     emit("aten::log.int : (int) -> (float)")
     emit("aten::add.float_int : (float, int) -> (float)", has_folder=True)
+    emit("aten::mul.float_int : (float, int) -> (float)", has_folder=True)
     emit("aten::sub.float : (float, float) -> (float)", has_folder=True)
     emit("aten::mul.float : (float, float) -> (float)", has_folder=True)
     emit("aten::div.float : (float, float) -> (float)", has_folder=True)

test/Conversion/TorchToArith/basic.mlir

@@ -236,6 +236,20 @@ func.func @torch.aten.mul.int(%arg0: !torch.int, %arg1: !torch.int) -> !torch.in
   return %0 : !torch.int
 }
 
+// CHECK-LABEL:  func.func @torch.aten.mul.float_int(
+// CHECK-SAME:                            %[[LHS:.*]]: !torch.float,
+// CHECK-SAME:                            %[[RHS:.*]]: !torch.int) -> !torch.float {
+// CHECK-DAG:      %[[LHS_F64:.*]] = torch_c.to_f64 %[[LHS]]
+// CHECK-DAG:      %[[RHS_I64:.*]] = torch_c.to_i64 %[[RHS]]
+// CHECK:          %[[RHS_F64:.*]] = arith.sitofp %[[RHS_I64]] : i64 to f64
+// CHECK:          %[[MUL:.*]] = arith.mulf %[[LHS_F64:.*]], [[RHS_F64:.*]] : f64
+// CHECK:          %[[OUT:.*]] = torch_c.from_f64 %[[MUL:.*]]
+// CHECK:          return %[[OUT:.*]] : !torch.float
+func.func @torch.aten.mul.float_int(%arg0: !torch.float, %arg1: !torch.int) -> !torch.float {
+  %0 = torch.aten.mul.float_int %arg0, %arg1 : !torch.float, !torch.int -> !torch.float
+  return %0 : !torch.float
+}
+
 // CHECK-LABEL:  func.func @torch.aten.div.float(
 // CHECK-SAME:                            %[[LHS:.*]]: !torch.float,
 // CHECK-SAME:                            %[[RHS:.*]]: !torch.float) -> !torch.float {

test/Dialect/Torch/canonicalize.mlir

@@ -1191,6 +1191,16 @@ func.func @torch.aten.mul.float() -> !torch.float {
     return %ret : !torch.float
 }
 
+// CHECK-LABEL:   func.func @torch.aten.mul.float_int() -> !torch.float {
+// CHECK:           %[[CST6:.*]] = torch.constant.float 6.000000e+00
+// CHECK:           return %[[CST6]] : !torch.float
+func.func @torch.aten.mul.float_int() -> !torch.float {
+    %cst2 = torch.constant.float 2.0
+    %cst3 = torch.constant.int 3
+    %ret = torch.aten.mul.float_int %cst2, %cst3: !torch.float, !torch.int -> !torch.float
+    return %ret : !torch.float
+}
+
 // CHECK-LABEL:   func.func @torch.aten.neg.float() -> !torch.float {
 // CHECK:           %[[CST_6:.*]] = torch.constant.float -6.000000e+00
 // CHECK:           return %[[CST_6]] : !torch.float