Skip to content

Commit cee8563

Browse files
afalkenberg1afalkenberg1
authored
[onnx] Support of onnx.Greater, onnx.Less, onnx.GreaterOrEqual to Torch (#2649)
The three remaining compare operations onnx.Greater onnx.Less onnx.GreaterOrEqual Are also added with this push request. This concludes a set of basic tensor compare functions.
1 parent 6188869 commit cee8563

File tree

2 files changed

+67
-4
lines changed

2 files changed

+67
-4
lines changed

‎lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp‎

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,43 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
3838
binder.op, resultType, lhs, rhs);
3939
return success();
4040
});
41-
patterns.onOp("LessOrEqual", 1,
41+
patterns.onOp("Greater", 16,
42+
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
43+
Torch::ValueTensorType resultType;
44+
Value lhs, rhs;
45+
std::string direction;
46+
if (binder.tensorOperands(lhs, rhs) ||
47+
binder.tensorResultType(resultType))
48+
return failure();
49+
rewriter.replaceOpWithNewOp<Torch::AtenGtTensorOp>(
50+
binder.op, resultType, lhs, rhs);
51+
return success();
52+
});
53+
patterns.onOp("GreaterOrEqual", 16,
54+
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
55+
Torch::ValueTensorType resultType;
56+
Value lhs, rhs;
57+
std::string direction;
58+
if (binder.tensorOperands(lhs, rhs) ||
59+
binder.tensorResultType(resultType))
60+
return failure();
61+
rewriter.replaceOpWithNewOp<Torch::AtenGeTensorOp>(
62+
binder.op, resultType, lhs, rhs);
63+
return success();
64+
});
65+
patterns.onOp("Less", 13,
66+
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
67+
Torch::ValueTensorType resultType;
68+
Value lhs, rhs;
69+
if (binder.tensorOperands(lhs, rhs) ||
70+
binder.tensorResultType(resultType)) {
71+
return failure();
72+
}
73+
rewriter.replaceOpWithNewOp<Torch::AtenLtTensorOp>(
74+
binder.op, resultType, lhs, rhs);
75+
return success();
76+
});
77+
patterns.onOp("LessOrEqual", 16,
4278
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
4379
Torch::ValueTensorType resultType;
4480
Value lhs, rhs;
@@ -87,4 +123,4 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
87123
binder.op, resultType, operand, constAlpha);
88124
return success();
89125
});
90-
}
126+
}

‎test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,33 @@
44
// level constants. This is a pragmatic choice which lets us have a lot
55
// of tests in this file, whereas the others tend to be more bespoke.
66

7+
// CHECK-LABEL: func.func @test_greater
8+
func.func @test_greater(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 16 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
9+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
10+
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
11+
// CHECK: torch.aten.gt.Tensor %[[ARG0]], %[[ARG1]] : !torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],i1>
12+
%0 = torch.operator "onnx.Greater"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1>
13+
return %0 : !torch.vtensor<[3,4,5],i1>
14+
}
15+
16+
// CHECK-LABEL: func.func @test_greater_or_equal
17+
func.func @test_greater_or_equal(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 16 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
18+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
19+
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
20+
// CHECK: torch.aten.ge.Tensor %[[ARG0]], %[[ARG1]] : !torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],i1>
21+
%0 = torch.operator "onnx.GreaterOrEqual"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1>
22+
return %0 : !torch.vtensor<[3,4,5],i1>
23+
}
24+
25+
// CHECK-LABEL: func.func @test_less
26+
func.func @test_less(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
27+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
28+
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
29+
// CHECK: torch.aten.lt.Tensor %[[ARG0]], %[[ARG1]] : !torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],i1>
30+
%0 = torch.operator "onnx.Less"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1>
31+
return %0 : !torch.vtensor<[3,4,5],i1>
32+
}
33+
734
// CHECK-LABEL: func.func @test_gather_elements
835
func.func @test_gather_elements(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5], si64>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.opset_version = 13 : si64} {
936
// CHECK-DAG: %[[INT0:.+]] = torch.constant.int 0
@@ -43,10 +70,10 @@ func.func @test_matmul_4d(%arg0: !torch.vtensor<[1,2,3,4],f32>, %arg1: !torch.vt
4370
}
4471

4572
// CHECK-LABEL: func.func @test_less_or_equal
46-
func.func @test_less_or_equal(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
73+
func.func @test_less_or_equal(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 16 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
4774
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
4875
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>
4976
// CHECK: torch.aten.le.Tensor %[[ARG0]], %[[ARG1]] : !torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],i1>
5077
%0 = torch.operator "onnx.LessOrEqual"(%arg0, %arg1) : (!torch.vtensor<[3,4,5],f32>, !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1>
5178
return %0 : !torch.vtensor<[3,4,5],i1>
52-
}
79+
}

0 commit comments

Comments
 (0)