llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[onnx] Support integer types for `onnx.Pow` (#3626) 

Pow is not support for the `torch` operator. Add casting for integer
types.
pull/3631/head
Rob Suderman 2024-08-13 09:39:04 -07:00 committed by GitHub
parent 39307f0462
commit af67f9efb0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 82 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
View File

@ -2856,16 +2856,64 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
binder.op, resultType, data, padsSizeList, modeVal, constantValue);
return success();
});
patterns.onOp("Pow", 1,
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
patterns.onOp(
"Pow", 1, [](OpBinder binder, ConversionPatternRewriter &rewriter) {
Torch::ValueTensorType resultType;
Value lhs, rhs;
if (binder.tensorOperands(lhs, rhs) ||
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(

21
test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir
View File

@ -1009,11 +1009,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 = ""} {
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>
}
// -----