[OnnxToTorch] Casting float to integer should round to nearest for pow with int result type (#4228)

Fixes #4091. ~~I assume this
will also need to be fixed for AtenPowScalarOp and AtenPowTensorScalarOp
as well. I'm putting up a PR to ensure the initial approach is correct
(new contributor :D ) before I put up another fix for AtenPowScalarOp
and AtenPowTensorScalarOp.~~

Author: cats-marin

lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp

@@ -3076,6 +3076,7 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
             rewriter.getIntegerAttr(rewriter.getIntegerType(64),
                                     static_cast<int64_t>(outDtype)));
 
+        pow = rewriter.create<Torch::AtenRoundOp>(loc, powType, pow);
         rewriter.replaceOpWithNewOp<Torch::AtenToDtypeOp>(
             binder.op, resultType, pow, outTyConst, cstFalse, cstFalse, none);
 

test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir

@@ -1253,14 +1253,30 @@ func.func @test_pow_i32(%arg0: !torch.vtensor<[3,4,5],si32>, %arg1: !torch.vtens
   // CHECK: %[[NONE:.+]] = torch.constant.none
   // CHECK: %[[POW:.+]] =  torch.aten.pow.Tensor_Tensor %arg0, %arg1 : !torch.vtensor<[3,4,5],si32>, !torch.vtensor<[3,4,5],si32> -> !torch.vtensor<[3,4,5],f64>
   // CHECK: %[[DTY:.+]] = torch.constant.int 3
-  // CHECK: %[[RES:.+]] = torch.aten.to.dtype %[[POW]], %[[DTY]], %[[FALSE]], %[[FALSE]], %[[NONE]]
+  // CHECK: %[[ROUND:.+]] = torch.aten.round %[[POW]] : !torch.vtensor<[3,4,5],f64> -> !torch.vtensor<[3,4,5],f64>
+  // CHECK: %[[RES:.+]] = torch.aten.to.dtype %[[ROUND]], %[[DTY]], %[[FALSE]], %[[FALSE]], %[[NONE]]
   // CHECK: return %[[RES]]
   %0 = torch.operator "onnx.Pow"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],si32>, !torch.vtensor<[3,4,5],si32>) -> !torch.vtensor<[3,4,5],si32>
   return %0 : !torch.vtensor<[3,4,5],si32>
 }
 
 // -----
 
+// CHECK-LABEL: func.func @test_pow_i32_f32_to_i32
+func.func @test_pow_i32_f32_to_i32(%arg0: !torch.vtensor<[3,4,5],si32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],si32> attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 15 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
+  // CHECK: %[[FALSE:.+]] = torch.constant.bool false
+  // CHECK: %[[NONE:.+]] = torch.constant.none
+  // CHECK: %[[POW:.+]] =  torch.aten.pow.Tensor_Tensor %arg0, %arg1 : !torch.vtensor<[3,4,5],si32>, !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],f64>
+  // CHECK: %[[DTY:.+]] = torch.constant.int 3
+  // CHECK: %[[ROUND:.+]] = torch.aten.round %[[POW]] : !torch.vtensor<[3,4,5],f64> -> !torch.vtensor<[3,4,5],f64>
+  // CHECK: %[[RES:.+]] = torch.aten.to.dtype %[[ROUND]], %[[DTY]], %[[FALSE]], %[[FALSE]], %[[NONE]]
+  // CHECK: return %[[RES]]
+  %0 = torch.operator "onnx.Pow"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],si32>, !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],si32>
+  return %0 : !torch.vtensor<[3,4,5],si32>
+}
+
+// -----
+
 // CHECK-LABEL: @test_hardsigmoid_example
 func.func @test_hardsigmoid_example(%arg0: !torch.vtensor<[3],f32>) -> !torch.vtensor<[3],f32> attributes {torch.onnx_meta.ir_version = 3 : si64, torch.onnx_meta.opset_version = 6 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
   // CHECK: %[[ALPHA_FLOAT:.*]] = torch.constant.float 5.000000e-01