xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/mhlo/canonicalize/reverse.mlir (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1// RUN: mlir-hlo-opt %s -pass-pipeline='func.func(canonicalize)' | FileCheck %s
2
3// CHECK-LABEL: func @noop
4// CHECK-SAME: (%[[ARG0:.*]]: tensor<1x2xf32>)
5func.func @noop(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
6  %0 = "mhlo.reverse"(%arg0) {dimensions = dense<[]> : tensor<0xi64>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
7  // CHECK: return %[[ARG0]]
8  func.return %0 : tensor<1x2xf32>
9}
10
11// CHECK-LABEL: func @dim1
12// CHECK-SAME: (%[[ARG0:.*]]: tensor
13func.func @dim1(%arg0: tensor<9x1x2x1x42xf32>) -> tensor<9x1x2x1x42xf32> {
14  %0 = "mhlo.reverse"(%arg0) {dimensions = dense<[1,3]> : tensor<2xi64>} : (tensor<9x1x2x1x42xf32>) -> tensor<9x1x2x1x42xf32>
15  // CHECK: return %[[ARG0]]
16  func.return %0 : tensor<9x1x2x1x42xf32>
17}
18
19// CHECK-LABEL: @noop_reverse_dynamic_shape
20// CHECK-SAME: [[ARG:%[a-zA-Z0-9]+]]
21func.func @noop_reverse_dynamic_shape(%arg0 : tensor<10x?x512xf32>) -> tensor<10x?x512xf32> {
22  %0 = "mhlo.reverse"(%arg0) {dimensions = dense<[0,1]> : tensor<2xi64>}: (tensor<10x?x512xf32>) -> tensor<10x?x512xf32>
23  // CHECK-NEXT: "mhlo.reverse"([[ARG]])
24  func.return %0 : tensor<10x?x512xf32>
25}
26
27// CHECK-LABEL: func @reverse_fold_constant_int
28func.func @reverse_fold_constant_int() -> tensor<0x2x0xi64> {
29  %cst = mhlo.constant dense<> : tensor<0x2x0xi64>
30  // CHECK: mhlo.constant dense<>
31  %1 = "mhlo.reverse"(%cst) {dimensions = dense<[0,1]> : tensor<2xi64>} : (tensor<0x2x0xi64>) -> tensor<0x2x0xi64>
32  func.return %1 : tensor<0x2x0xi64>
33}
34
35// CHECK-LABEL: func @reverse_fold_constant_int_0
36func.func @reverse_fold_constant_int_0() -> tensor<0xi64> {
37  %cst = mhlo.constant dense<> : tensor<0xi64>
38  // CHECK: mhlo.constant dense<>
39  %1 = "mhlo.reverse"(%cst) {dimensions = dense<[0]> : tensor<1xi64>} : (tensor<0xi64>) -> tensor<0xi64>
40  func.return %1 : tensor<0xi64>
41}
42
43// CHECK-LABEL: func @reverse_fold_constant_int_1
44func.func @reverse_fold_constant_int_1() -> tensor<3x2xi32> {
45  %cst = mhlo.constant dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>
46  // CHECK: mhlo.constant dense<{{\[\[}}6, 5], [4, 3], [2, 1]]>
47  %1 = "mhlo.reverse"(%cst) {dimensions = dense<[0,1]> : tensor<2xi64>} : (tensor<3x2xi32>) -> tensor<3x2xi32>
48  func.return %1 : tensor<3x2xi32>
49}
50
51// CHECK-LABEL: func @reverse_fold_constant_int_2
52func.func @reverse_fold_constant_int_2() -> tensor<3x2xi32> {
53  %cst = mhlo.constant dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>
54  // CHECK: mhlo.constant dense<{{\[\[}}5, 6], [3, 4], [1, 2]]>
55  %1 = "mhlo.reverse"(%cst) {dimensions = dense<[0]> : tensor<1xi64>} : (tensor<3x2xi32>) -> tensor<3x2xi32>
56  func.return %1 : tensor<3x2xi32>
57}
58
59// CHECK-LABEL: func @reverse_fold_constant_int_3
60func.func @reverse_fold_constant_int_3() -> tensor<3x2xi32> {
61  %cst = mhlo.constant dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>
62  // CHECK: mhlo.constant dense<{{\[\[}}2, 1], [4, 3], [6, 5]]>
63  %1 = "mhlo.reverse"(%cst) {dimensions = dense<[1]> : tensor<1xi64>} : (tensor<3x2xi32>) -> tensor<3x2xi32>
64  func.return %1 : tensor<3x2xi32>
65}
66
67// CHECK-LABEL: func @reverse_fold_constant_int_4
68func.func @reverse_fold_constant_int_4() -> tensor<2x3x2xi32> {
69  %cst = mhlo.constant dense<[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]> : tensor<2x3x2xi32>
70  // CHECK: mhlo.constant dense<{{\[\[\[}}12, 11], [10, 9], [8, 7]], {{\[\[}}6, 5], [4, 3], [2, 1]]]>
71  %1 = "mhlo.reverse"(%cst) {dimensions = dense<[0,1,2]> : tensor<3xi64>} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
72  func.return %1 : tensor<2x3x2xi32>
73}
74
75// CHECK-LABEL: func @reverse_fold_constant_float
76func.func @reverse_fold_constant_float() -> tensor<3x2xf32> {
77  %cst = mhlo.constant dense<[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]> : tensor<3x2xf32>
78  // CHECK: mhlo.constant dense<{{\[\[}}6.000000e+00, 5.000000e+00], [4.000000e+00, 3.000000e+00], [2.000000e+00, 1.000000e+00]]>
79  %1 = "mhlo.reverse"(%cst) {dimensions = dense<[0,1]> : tensor<2xi64>} : (tensor<3x2xf32>) -> tensor<3x2xf32>
80  func.return %1 : tensor<3x2xf32>
81}
82
83