[onnx] Support integer types for onnx.Pow

Pow is not support for the `torch` operator. Add casting for integer types.
llvm · Aug 9, 2024 · 3748a68 · 3748a68
1 parent 8358e8c
commit 3748a68
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 7 deletions.
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -2480,8 +2480,60 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
                       binder.tensorResultType(resultType)) {
                     return failure();
                   }
-                  rewriter.replaceOpWithNewOp<Torch::AtenPowTensorTensorOp>(
-                      binder.op, resultType, lhs, rhs);
+
+                  auto loc = binder.getLoc();
+                  auto lhsTy = cast<Torch::ValueTensorType>(lhs.getType());
+                  auto rhsTy = cast<Torch::ValueTensorType>(rhs.getType());
+                  Value cstFalse = rewriter.create<Torch::ConstantBoolOp>(
+                      loc, rewriter.getBoolAttr(false));
+                  Value none = rewriter.create<Torch::ConstantNoneOp>(loc);
+                  auto torchDtype =
+                      Torch::getScalarTypeForType(rewriter.getF32Type());
+                  Value tyConst = rewriter.create<Torch::ConstantIntOp>(
+                      binder.getLoc(), rewriter.getType<Torch::IntType>(),
+                      rewriter.getIntegerAttr(
+                          rewriter.getIntegerType(64),
+                          static_cast<int64_t>(torchDtype)));
+
+                  if (isa<IntegerType>(lhsTy.getDtype())) {
+                    lhsTy = rewriter.getType<Torch::ValueTensorType>(
+                        lhsTy.getSizes(), rewriter.getF32Type());
+                    lhs = rewriter.create<Torch::AtenToDtypeOp>(
+                        loc, lhsTy, lhs, tyConst, cstFalse, cstFalse, none);
+                  }
+
+                  if (isa<IntegerType>(rhsTy.getDtype())) {
+                    rhsTy = rewriter.getType<Torch::ValueTensorType>(
+                        rhsTy.getSizes(), rewriter.getF32Type());
+                    rhs = rewriter.create<Torch::AtenToDtypeOp>(
+                        loc, rhsTy, rhs, tyConst, cstFalse, cstFalse, none);
+                  }
+
+                  auto powType = resultType;
+                  if (isa<IntegerType>(resultType.getDtype())) {
+                    powType = rewriter.getType<Torch::ValueTensorType>(
+                        resultType.getSizes(), rewriter.getF32Type());
+                  }
+
+                  Value pow = rewriter.create<Torch::AtenPowTensorTensorOp>(
+                      loc, powType, lhs, rhs);
+
+                  if (!isa<IntegerType>(resultType.getDtype())) {
+                    rewriter.replaceOp(binder.op, pow);
+                    return success();
+                  }
+
+                  auto outDtype =
+                      Torch::getScalarTypeForType(resultType.getDtype());
+                  auto outTyConst = rewriter.create<Torch::ConstantIntOp>(
+                      binder.getLoc(), rewriter.getType<Torch::IntType>(),
+                      rewriter.getIntegerAttr(rewriter.getIntegerType(64),
+                                              static_cast<int64_t>(outDtype)));
+
+                  rewriter.replaceOpWithNewOp<Torch::AtenToDtypeOp>(
+                      binder.op, resultType, pow, outTyConst, cstFalse,
+                      cstFalse, none);
+
                   return success();
                 });
   patterns.onOp(

diff --git a/test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir b/test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir
@@ -974,11 +974,28 @@ func.func @test_pad_edge(%arg0: !torch.vtensor<[3,4],f32>, %arg1: !torch.vtensor
 // -----
 
 // CHECK-LABEL: func.func @test_pow
-  func.func @test_pow(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> 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: torch.aten.pow.Tensor_Tensor %arg0, %arg1 : !torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],f32>
-    %0 = torch.operator "onnx.Pow"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32>
-    return %0 : !torch.vtensor<[3,4,5],f32>
-  }
+func.func @test_pow(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> 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: torch.aten.pow.Tensor_Tensor %arg0, %arg1 : !torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],f32>
+  %0 = torch.operator "onnx.Pow"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32>
+  return %0 : !torch.vtensor<[3,4,5],f32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @test_pow_i32
+func.func @test_pow_i32(%arg0: !torch.vtensor<[3,4,5],si32>, %arg1: !torch.vtensor<[3,4,5],si32>) -> !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: %[[DTY:.+]] = torch.constant.int 6
+  // CHECK: %[[CAST_LHS:.+]] = torch.aten.to.dtype %arg0, %[[DTY]], %[[FALSE]], %[[FALSE]], %[[NONE]]
+  // CHECK: %[[CAST_RHS:.+]] = torch.aten.to.dtype %arg1, %[[DTY]], %[[FALSE]], %[[FALSE]], %[[NONE]]
+  // CHECK: %[[POW:.+]] = torch.aten.pow.Tensor_Tensor %[[CAST_LHS]], %[[CAST_RHS]]
+  // CHECK: %[[DTY:.+]] = torch.constant.int 3
+  // CHECK: %[[RES:.+]] = torch.aten.to.dtype %2, %[[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>
+}
 
 // -----