| // RUN: mlir-hlo-opt %s -pass-pipeline='func(canonicalize)' | FileCheck %s |
| |
| // CHECK-LABEL: add_fold |
| func @add_fold() -> tensor<4xi64> { |
| %0 = mhlo.constant dense<[1, 2, 3, 4]> : tensor<4xi64> |
| %1 = mhlo.constant dense<[5, 6, 7, 8]> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<[6, 8, 10, 12]> |
| %2 = "mhlo.add"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) |
| return %2 : tensor<4xi64> |
| } |
| |
| // CHECK-LABEL: add_scalar_fold |
| func @add_scalar_fold() -> tensor<4xi64> { |
| %0 = mhlo.constant dense<1> : tensor<4xi64> |
| %1 = mhlo.constant dense<5> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<6> |
| %2 = "mhlo.add"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) |
| return %2 : tensor<4xi64> |
| } |
| |
| // CHECK-LABEL: add_fold_float |
| func @add_fold_float() -> tensor<4xf64> { |
| %0 = mhlo.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf64> |
| %1 = mhlo.constant dense<[5.0, 6.0, 7.0, 8.0]> : tensor<4xf64> |
| // CHECK: mhlo.constant dense<[6.000000e+00, 8.000000e+00, 1.000000e+01, 1.200000e+01]> |
| %2 = "mhlo.add"(%0, %1) : (tensor<4xf64>, tensor<4xf64>) -> (tensor<4xf64>) |
| return %2 : tensor<4xf64> |
| } |
| |
| // CHECK-LABEL: sub_scalar_fold |
| func @sub_scalar_fold() -> tensor<4xi64> { |
| %0 = mhlo.constant dense<5> : tensor<4xi64> |
| %1 = mhlo.constant dense<1> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<4> |
| %2 = "mhlo.subtract"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) |
| return %2 : tensor<4xi64> |
| } |
| |
| // CHECK-LABEL: multiply_scalar_fold |
| func @multiply_scalar_fold() -> tensor<4xi64> { |
| %0 = mhlo.constant dense<5> : tensor<4xi64> |
| %1 = mhlo.constant dense<3> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<15> |
| %2 = "mhlo.multiply"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) |
| return %2 : tensor<4xi64> |
| } |
| |
| // CHECK-LABEL: divide_scalar_fold |
| func @divide_scalar_fold() -> tensor<4xi64> { |
| %0 = mhlo.constant dense<7> : tensor<4xi64> |
| %1 = mhlo.constant dense<5> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<1> |
| %2 = "mhlo.divide"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) |
| return %2 : tensor<4xi64> |
| } |
| |
| // CHECK-LABEL: divide_fold_float |
| func @divide_fold_float() -> tensor<4xf64> { |
| %0 = mhlo.constant dense<[5.0, 66.0, 5.0, 1.0]> : tensor<4xf64> |
| %1 = mhlo.constant dense<[5.0, 3.0, 2.0, 4.0]> : tensor<4xf64> |
| // CHECK: mhlo.constant dense<[1.000000e+00, 2.200000e+01, 2.500000e+00, 2.500000e-01]> |
| %2 = "mhlo.divide"(%0, %1) : (tensor<4xf64>, tensor<4xf64>) -> (tensor<4xf64>) |
| return %2 : tensor<4xf64> |
| } |
| |
| // CHECK-LABEL: remainder_fold_int |
| func @remainder_fold_int() -> tensor<4xi32> { |
| %0 = mhlo.constant dense<[5, 66, 5, 1]> : tensor<4xi32> |
| %1 = mhlo.constant dense<[3, 5, 1, 2]> : tensor<4xi32> |
| // CHECK: mhlo.constant dense<[2, 1, 0, 1]> |
| %2 = "mhlo.remainder"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> (tensor<4xi32>) |
| return %2 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: remainder_fold_float |
| func @remainder_fold_float() -> tensor<4xf32> { |
| %0 = mhlo.constant dense<[7.0, 66.5, 5.0, 3.1]> : tensor<4xf32> |
| %1 = mhlo.constant dense<[3.0, 5.0, 1.0, 2.6]> : tensor<4xf32> |
| // CHECK: mhlo.constant dense<[1.000000e+00, 1.500000e+00, 0.000000e+00, 5.000000e-01]> |
| %2 = "mhlo.remainder"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> (tensor<4xf32>) |
| return %2 : tensor<4xf32> |
| } |
| |
| // CHECK-LABEL: round_fold |
| func @round_fold() -> tensor<4xf32> { |
| %0 = mhlo.constant dense<[-1.5, -0.1, 1.1, 2.5]> : tensor<4xf32> |
| %1 = "mhlo.round_nearest_afz"(%0) : (tensor<4xf32>) -> tensor<4xf32> |
| return %1 : tensor<4xf32> |
| // CHECK: mhlo.constant dense<[-2.000000e+00, -0.000000e+00, 1.000000e+00, 3.000000e+00]> |
| } |
| |
| // CHECK-LABEL: max_scalar_fold |
| func @max_scalar_fold() -> tensor<4xi64> { |
| %0 = mhlo.constant dense<7> : tensor<4xi64> |
| %1 = mhlo.constant dense<5> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<7> |
| %2 = "mhlo.maximum"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) |
| return %2 : tensor<4xi64> |
| } |
| |
| // CHECK-LABEL: max_fold_float |
| func @max_fold_float() -> tensor<4xf64> { |
| %0 = mhlo.constant dense<[5.0, 66.0, 5.0, 1.0]> : tensor<4xf64> |
| %1 = mhlo.constant dense<[5.0, 3.0, 2.0, 4.0]> : tensor<4xf64> |
| // CHECK: mhlo.constant dense<[5.000000e+00, 6.600000e+01, 5.000000e+00, 4.000000e+00]> |
| %2 = "mhlo.maximum"(%0, %1) : (tensor<4xf64>, tensor<4xf64>) -> (tensor<4xf64>) |
| return %2 : tensor<4xf64> |
| } |
| |
| // CHECK-LABEL: min_scalar_fold |
| func @min_scalar_fold() -> tensor<4xi64> { |
| %0 = mhlo.constant dense<7> : tensor<4xi64> |
| %1 = mhlo.constant dense<-5> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<-5> |
| %2 = "mhlo.minimum"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) |
| return %2 : tensor<4xi64> |
| } |
| |
| // CHECK-LABEL: min_fold_float |
| func @min_fold_float() -> tensor<4xf64> { |
| %0 = mhlo.constant dense<[5.0, 66.0, 5.0, 1.0]> : tensor<4xf64> |
| %1 = mhlo.constant dense<[5.0, 3.0, 2.0, 4.0]> : tensor<4xf64> |
| // CHECK: mhlo.constant dense<[5.000000e+00, 3.000000e+00, 2.000000e+00, 1.000000e+00]> |
| %2 = "mhlo.minimum"(%0, %1) : (tensor<4xf64>, tensor<4xf64>) -> (tensor<4xf64>) |
| return %2 : tensor<4xf64> |
| } |
| |
| // CHECK-LABEL: concatenate_noop |
| func @concatenate_noop(%arg0: tensor<4xi32>) -> tensor<4xi32> { |
| // CHECK-SAME: [[ARG:%.+]]: tensor<4xi32> |
| %0 = "mhlo.concatenate"(%arg0) { dimension = 0 : i64 } : (tensor<4xi32>) -> tensor<4xi32> |
| |
| // CHECK: return [[ARG]] |
| return %0 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: concatenate_remove_operand |
| func @concatenate_remove_operand(%arg0: tensor<4xi32>, %arg1: tensor<0xi32>) -> tensor<4xi32> { |
| // CHECK-SAME: [[ARG0:%.+]]: tensor<4xi32> |
| // CHECK-SAME: [[ARG1:%.+]]: tensor<0xi32> |
| %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<4xi32>, tensor<0xi32>) -> tensor<4xi32> |
| |
| // CHECK: return [[ARG0]] |
| return %0 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: concatenate_empty_bool |
| func @concatenate_empty_bool(%arg0: tensor<0xi1>, %arg1: tensor<0xi1>) -> tensor<0xi1> { |
| // CHECK: mhlo.constant |
| %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<0xi1>, tensor<0xi1>) -> tensor<0xi1> |
| |
| return %0 : tensor<0xi1> |
| } |
| |
| // CHECK-LABEL: concatenate_empty_int |
| func @concatenate_empty_int(%arg0: tensor<0xi32>, %arg1: tensor<0xi32>) -> tensor<0xi32> { |
| // CHECK: mhlo.constant |
| %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi32> |
| |
| return %0 : tensor<0xi32> |
| } |
| |
| // CHECK-LABEL: concatenate_empty_float |
| func @concatenate_empty_float(%arg0: tensor<0xf32>, %arg1: tensor<0xf32>) -> tensor<0xf32> { |
| // CHECK: mhlo.constant |
| %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<0xf32>, tensor<0xf32>) -> tensor<0xf32> |
| |
| return %0 : tensor<0xf32> |
| } |
| |
| // CHECK-LABEL: concatenate_const_1D |
| func @concatenate_const_1D() -> tensor<4xi32> { |
| // CHECK: [[VAL:%.+]]= mhlo.constant dense<[0, 1, 2, 3]> |
| %0 = mhlo.constant dense<[0, 1]> : tensor<2xi32> |
| %1 = mhlo.constant dense<[2, 3]> : tensor<2xi32> |
| %2 = "mhlo.concatenate"(%0, %1) { dimension = 0 : i64 } : (tensor<2xi32>, tensor<2xi32>) -> tensor<4xi32> |
| |
| // CHECK: return [[VAL]] |
| return %2 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: concatenate_const_1D_float |
| func @concatenate_const_1D_float() -> tensor<4xf32> { |
| // CHECK: [[VAL:%.+]] = mhlo.constant dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]> |
| |
| %0 = mhlo.constant dense<[0.0, 1.0]> : tensor<2xf32> |
| %1 = mhlo.constant dense<[2.0, 3.0]> : tensor<2xf32> |
| %2 = "mhlo.concatenate"(%0, %1) { dimension = 0 : i64 } : (tensor<2xf32>, tensor<2xf32>) -> tensor<4xf32> |
| |
| // CHECK: return [[VAL]] |
| return %2 : tensor<4xf32> |
| } |
| |
| // CHECK-LABEL: concatenate_const_2D_vertical |
| func @concatenate_const_2D_vertical() -> tensor<2x2xi32> { |
| // CHECK: [[VAL:%.+]]= mhlo.constant dense<[ |
| // CHECK-SAME: [0, 1], [2, 3] |
| // CHECK-SAME: ]> |
| %0 = mhlo.constant dense<[[0, 1]]> : tensor<1x2xi32> |
| %1 = mhlo.constant dense<[[2, 3]]> : tensor<1x2xi32> |
| %2 = "mhlo.concatenate"(%0, %1) { dimension = 0 : i64 } : (tensor<1x2xi32>, tensor<1x2xi32>) -> tensor<2x2xi32> |
| |
| // CHECK: return [[VAL]] |
| return %2 : tensor<2x2xi32> |
| } |
| |
| // CHECK-LABEL: concatenate_const_2D_horizontal |
| func @concatenate_const_2D_horizontal() -> tensor<2x2xi32> { |
| // CHECK: [[VAL:%.+]]= mhlo.constant dense<[ |
| // CHECK-SAME: [0, 2], [1, 3] |
| // CHECK-SAME: ]> |
| %0 = mhlo.constant dense<[[0], [1]]> : tensor<2x1xi32> |
| %1 = mhlo.constant dense<[[2], [3]]> : tensor<2x1xi32> |
| %2 = "mhlo.concatenate"(%0, %1) { dimension = 1 : i64 } : (tensor<2x1xi32>, tensor<2x1xi32>) -> tensor<2x2xi32> |
| |
| // CHECK: return [[VAL]] |
| return %2 : tensor<2x2xi32> |
| } |
| |
| // CHECK-LABEL: constant_like_constant |
| func @constant_like_constant(%arg0: tensor<3x4xi32>) -> tensor<3x4xf32> { |
| // CHECK: mhlo.constant dense<3.200000e+00> |
| %0 = "chlo.constant_like"(%arg0) { value = 3.2 : f32 } : (tensor<3x4xi32>) -> tensor<3x4xf32> |
| return %0 : tensor<3x4xf32> |
| } |
| |
| // CHECK-LABEL: constant_like_constant_dynamic |
| func @constant_like_constant_dynamic(%arg0: tensor<*xi32>) -> tensor<*xf32> { |
| // CHECK: chlo.constant_like |
| %0 = "chlo.constant_like"(%arg0) { value = 3.2 : f32 } : (tensor<*xi32>) -> tensor<*xf32> |
| return %0 : tensor<*xf32> |
| } |
| |
| // CHECK-LABEL: dynamic_slice_variable_start |
| func @dynamic_slice_variable_start(%arg0: tensor<3x4xi32>, %arg1: tensor<i64>, %arg2: tensor<i64>) -> tensor<1x4xi32> { |
| // CHECK: "mhlo.dynamic-slice" |
| %1 = "mhlo.dynamic-slice"(%arg0, %arg1, %arg2) {slice_sizes = dense<[1, 4]> : tensor<2xi64>} : (tensor<3x4xi32>, tensor<i64>, tensor<i64>) -> tensor<1x4xi32> |
| return %1 : tensor<1x4xi32> |
| } |
| |
| // CHECK-LABEL: dynamic_slice_constant_start |
| func @dynamic_slice_constant_start(%arg0: tensor<4xi32>) -> tensor<2xi32> { |
| // CHECK: %[[RESULT:.*]] = "mhlo.slice"(%arg0) |
| // CHECK-DAG-SAME: limit_indices = dense<3> : tensor<1xi64> |
| // CHECK-DAG-SAME: start_indices = dense<1> : tensor<1xi64> |
| // CHECK-DAG-SAME: strides = dense<1> : tensor<1xi64>} |
| // CHECK: return %[[RESULT]] : tensor<2xi32> |
| %0 = mhlo.constant dense<1> : tensor<i64> |
| %1 = "mhlo.dynamic-slice"(%arg0, %0) {slice_sizes = dense<2> : tensor<1xi64>} : (tensor<4xi32>, tensor<i64>) -> tensor<2xi32> |
| return %1 : tensor<2xi32> |
| } |
| |
| // CHECK-LABEL: dynamic_slice_constant_start_dynamic_shape |
| func @dynamic_slice_constant_start_dynamic_shape(%arg0: tensor<?x4xi32>, %arg1: tensor<2xi64>) -> tensor<?x4xi32> { |
| // CHECK: %[[RESULT:.*]] = "mhlo.slice"(%arg0) |
| // CHECK-DAG-SAME: limit_indices = dense<[2, 4]> : tensor<2xi64> |
| // CHECK-DAG-SAME: start_indices = dense<[1, 0]> : tensor<2xi64> |
| // CHECK-DAG-SAME: strides = dense<1> : tensor<2xi64> |
| // CHECK: return %[[RESULT]] : tensor<?x4xi32> |
| %0 = mhlo.constant dense<1> : tensor<i64> |
| %1 = mhlo.constant dense<0> : tensor<i64> |
| %2 = "mhlo.dynamic-slice"(%arg0, %0, %1) {slice_sizes = dense<[1, 4]> : tensor<2xi64>} : (tensor<?x4xi32>, tensor<i64>, tensor<i64>) -> tensor<?x4xi32> |
| return %2 : tensor<?x4xi32> |
| } |
| |
| // CHECK-LABEL: slice_2D_noop |
| // CHECK-SAME: [[ARG:%.+]]: tensor<2x2xi64> |
| func @slice_2D_noop(%arg0: tensor<2x2xi64>) -> tensor<2x2xi64> { |
| %0 = "mhlo.slice"(%arg0) { limit_indices = dense<[2, 2]> : tensor<2xi64>, start_indices = dense<[0, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xi64>) -> (tensor<2x2xi64>) |
| |
| // CHECK-NEXT: return [[ARG]] |
| return %0 : tensor<2x2xi64> |
| } |
| |
| // CHECK-LABEL: slice_1D_fold |
| func @slice_1D_fold() -> tensor<2xi64> { |
| %0 = mhlo.constant dense<[5, 7, 9, 10]> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<[7, 9]> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[3]> : tensor<1xi64>, start_indices = dense<[1]> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xi64>) -> (tensor<2xi64>) |
| return %1 : tensor<2xi64> |
| } |
| |
| // CHECK-LABEL: slice_1D_fp |
| func @slice_1D_fp() -> tensor<2xf32> { |
| %0 = mhlo.constant dense<[5.0, 7.0, 9.0, 10.0]> : tensor<4xf32> |
| // CHECK: mhlo.constant dense<[7.000000e+00, 9.000000e+00]> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[3]> : tensor<1xi64>, start_indices = dense<[1]> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xf32>) -> (tensor<2xf32>) |
| return %1 : tensor<2xf32> |
| } |
| |
| // CHECK-LABEL: slice_1D_strided_fold |
| func @slice_1D_strided_fold() -> tensor<2xi64> { |
| %0 = mhlo.constant dense<[5, 7, 9, 10]> : tensor<4xi64> |
| // CHECK: mhlo.constant dense<[7, 10]> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[4]> : tensor<1xi64>, start_indices = dense<[1]> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>} : (tensor<4xi64>) -> (tensor<2xi64>) |
| return %1 : tensor<2xi64> |
| } |
| |
| // CHECK-LABEL: slice_2D_fold |
| func @slice_2D_fold() -> tensor<2x2xi64> { |
| %0 = mhlo.constant dense<[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]> : tensor<4x4xi64> |
| // CHECK-NEXT: mhlo.constant dense<[ |
| // CHECK-SAME: [6, 7], |
| // CHECK-SAME: [10, 11] |
| // CHECK-SAME: ]> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[3, 4]> : tensor<2xi64>, start_indices = dense<[1, 2]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<4x4xi64>) -> (tensor<2x2xi64>) |
| return %1 : tensor<2x2xi64> |
| } |
| |
| // CHECK-LABEL: slice_2D_fold_horizontal |
| func @slice_2D_fold_horizontal() -> tensor<1x4xi64> { |
| %0 = mhlo.constant dense<[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]> : tensor<4x4xi64> |
| // CHECK-NEXT: mhlo.constant dense<[ |
| // CHECK-SAME: [0, 1, 2, 3] |
| // CHECK-SAME: ]> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[1, 4]> : tensor<2xi64>, start_indices = dense<[0, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<4x4xi64>) -> (tensor<1x4xi64>) |
| return %1 : tensor<1x4xi64> |
| } |
| |
| // CHECK-LABEL: slice_2D_fold_vertical |
| func @slice_2D_fold_vertical() -> tensor<4x1xi64> { |
| %0 = mhlo.constant dense<[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]> : tensor<4x4xi64> |
| // CHECK-NEXT: mhlo.constant dense<[ |
| // CHECK-SAME: [2], [6], [10], [14] |
| // CHECK-SAME: ]> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[4, 3]> : tensor<2xi64>, start_indices = dense<[0, 2]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<4x4xi64>) -> (tensor<4x1xi64>) |
| return %1 : tensor<4x1xi64> |
| } |
| |
| // CHECK-LABEL: slice_zero_elements |
| func @slice_zero_elements() -> tensor<0xi64> { |
| %0 = mhlo.constant dense<> : tensor<0xi64> |
| // CHECK: %[[CONST:.*]] = mhlo.constant dense<> : tensor<0xi64> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[0]> : tensor<1xi64>, start_indices = dense<[0]> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<0xi64>) -> (tensor<0xi64>) |
| // CHECK: return %[[CONST]] : tensor<0xi64> |
| return %1 : tensor<0xi64> |
| } |
| |
| // CHECK-LABEL: slice_unknown_shape |
| func @slice_unknown_shape(%arg0: tensor<*xf32>) -> tensor<*xf32> { |
| // CHECK: "mhlo.slice"(%arg0) {limit_indices = dense<[1, 4]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<*xf32>) -> tensor<*xf32> |
| %0 = "mhlo.slice"(%arg0) {limit_indices = dense<[1, 4]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<*xf32>) -> tensor<*xf32> |
| return %0 : tensor<*xf32> |
| } |
| |
| // CHECK-LABEL: slice_concat_fold_first |
| func @slice_concat_fold_first(%arg0: tensor<1x5xf32>, %arg1: tensor<1x5xf32>) -> tensor<1x5xf32> { |
| %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<1x5xf32>, tensor<1x5xf32>) -> tensor<2x5xf32> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[1, 5]> : tensor<2xi64>, start_indices = dense<[0, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x5xf32>) -> (tensor<1x5xf32>) |
| // CHECK: return %arg0 |
| return %1 : tensor<1x5xf32> |
| } |
| |
| // CHECK-LABEL: slice_concat_fold_second |
| func @slice_concat_fold_second(%arg0: tensor<1x5xf32>, %arg1: tensor<1x5xf32>) -> tensor<1x5xf32> { |
| %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<1x5xf32>, tensor<1x5xf32>) -> tensor<2x5xf32> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[2, 5]> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x5xf32>) -> (tensor<1x5xf32>) |
| // CHECK: return %arg1 |
| return %1 : tensor<1x5xf32> |
| } |
| |
| // CHECK-LABEL: slice_concat_fold_second_with_slice |
| func @slice_concat_fold_second_with_slice(%arg0: tensor<1x5xf32>, %arg1: tensor<1x5xf32>) -> tensor<1x4xf32> { |
| %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<1x5xf32>, tensor<1x5xf32>) -> tensor<2x5xf32> |
| // CHECK: [[SLICE:%.+]] = "mhlo.slice"(%arg1) {limit_indices = dense<[1, 5]> : tensor<2xi64>, start_indices = dense<[0, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x5xf32>) -> tensor<1x4xf32> |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[2, 5]> : tensor<2xi64>, start_indices = dense<[1, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x5xf32>) -> (tensor<1x4xf32>) |
| |
| // CHECK: return [[SLICE]] |
| return %1 : tensor<1x4xf32> |
| } |
| |
| // CHECK-LABEL: slice_concat_fold_middle |
| func @slice_concat_fold_middle(%arg0: tensor<1x5xf32>, %arg1: tensor<2x5xf32>, %arg2: tensor<1x5xf32>) -> tensor<1x5xf32> { |
| %0 = "mhlo.concatenate"(%arg0, %arg1, %arg2) { dimension = 0 : i64 } : (tensor<1x5xf32>, tensor<2x5xf32>, tensor<1x5xf32>) -> tensor<4x5xf32> |
| // CHECK: [[SLICE:%.+]] = "mhlo.slice"(%arg1) {limit_indices = dense<[2, 5]> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[3, 5]> : tensor<2xi64>, start_indices = dense<[2, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<4x5xf32>) -> (tensor<1x5xf32>) |
| |
| // CHECK: return [[SLICE]] |
| return %1 : tensor<1x5xf32> |
| } |
| |
| // CHECK-LABEL: slice_concat_fold_two |
| func @slice_concat_fold_two(%arg0: tensor<1x5xf32>, %arg1: tensor<2x5xf32>, %arg2: tensor<1x5xf32>) -> tensor<2x5xf32> { |
| // CHECK: [[CONCAT:%.+]] = "mhlo.concatenate"(%arg1, %arg2) {dimension = 0 : i64} |
| %0 = "mhlo.concatenate"(%arg0, %arg1, %arg2) { dimension = 0 : i64 } : (tensor<1x5xf32>, tensor<2x5xf32>, tensor<1x5xf32>) -> tensor<4x5xf32> |
| |
| // CHECK: [[SLICE:%.+]] = "mhlo.slice"([[CONCAT]]) {limit_indices = dense<[3, 5]> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} |
| %1 = "mhlo.slice"(%0) { limit_indices = dense<[4, 5]> : tensor<2xi64>, start_indices = dense<[2, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<4x5xf32>) -> (tensor<2x5xf32>) |
| |
| // CHECK: return [[SLICE]] |
| return %1 : tensor<2x5xf32> |
| } |
| |
| // CHECK-LABEL: func @broadcast_in_dim_identity |
| func @broadcast_in_dim_identity(%arg0: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> { |
| // CHECK: return %arg0 |
| %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>} : (tensor<2x3x4xf32>) -> tensor<2x3x4xf32> |
| return %0 : tensor<2x3x4xf32> |
| } |
| |
| // CHECK-LABEL: func @broadcast_in_dim_not_identity_because_it_actually_broadcasts |
| func @broadcast_in_dim_not_identity_because_it_actually_broadcasts(%arg0: tensor<1x2xf32>) -> tensor<2x2xf32> { |
| // CHECK: mhlo.broadcast_in_dim |
| %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x2xf32>) -> tensor<2x2xf32> |
| return %0 : tensor<2x2xf32> |
| } |
| |
| // CHECK-LABEL: func @broadcast_in_dim_not_identity_permutation |
| func @broadcast_in_dim_not_identity_permutation(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> { |
| // CHECK: mhlo.broadcast_in_dim |
| %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[1, 0]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<2x2xf32> |
| return %0 : tensor<2x2xf32> |
| } |
| |
| |
| // CHECK-LABEL: func @dynamic_broadcast_in_dim_op_not_actually_dynamic |
| func @dynamic_broadcast_in_dim_op_not_actually_dynamic(%arg0: tensor<4xf32>, %arg1: tensor<2xi64>) -> tensor<5x4xf32> { |
| // CHECK: %[[RESULT:.+]] = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<4xf32>) -> tensor<5x4xf32> |
| %0 = "mhlo.dynamic_broadcast_in_dim"(%arg0, %arg1) { broadcast_dimensions = dense<1> : tensor<1xi64> } : (tensor<4xf32>, tensor<2xi64>) -> tensor<5x4xf32> |
| // CHECK: return %[[RESULT]] : tensor<5x4xf32> |
| return %0 : tensor<5x4xf32> |
| } |
| |
| // CHECK-LABEL: func @dynamic_broadcast_in_dim_to_same_shape_1 |
| func @dynamic_broadcast_in_dim_to_same_shape_1(%arg0: tensor<?xf32>) -> tensor<?xf32> { |
| // CHECK-SAME: %[[ARG:.*]]: tensor<?xf32> |
| %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<1xindex> |
| %2 = "mhlo.dynamic_broadcast_in_dim"(%arg0, %0) { broadcast_dimensions = dense<0> : tensor<1xi64> } : (tensor<?xf32>, tensor<1xindex>) -> tensor<?xf32> |
| // CHECK: return %[[ARG]] : tensor<?xf32> |
| return %2 : tensor<?xf32> |
| } |
| |
| // CHECK-LABEL: func @dynamic_broadcast_in_dim_to_same_shape_2 |
| func @dynamic_broadcast_in_dim_to_same_shape_2(%arg0: tensor<?xf32>) -> tensor<?xf32> { |
| // CHECK-SAME: %[[ARG:.*]]: tensor<?xf32> |
| %0 = shape.shape_of %arg0 : tensor<?xf32> -> !shape.shape |
| %1 = shape.to_extent_tensor %0 : !shape.shape -> tensor<1xindex> |
| %2 = "mhlo.dynamic_broadcast_in_dim"(%arg0, %1) { broadcast_dimensions = dense<0> : tensor<1xi64> } : (tensor<?xf32>, tensor<1xindex>) -> tensor<?xf32> |
| // CHECK: return %[[ARG]] : tensor<?xf32> |
| return %2 : tensor<?xf32> |
| } |
| |
| // CHECK-LABEL: func @dynamic_broadcast_in_dim_to_same_shape_3 |
| func @dynamic_broadcast_in_dim_to_same_shape_3(%arg0: tensor<*xf32>) -> tensor<?xf32> { |
| // CHECK-SAME: %[[ARG:.*]]: tensor<*xf32> |
| %0 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex> |
| %1 = tensor.cast %0 : tensor<?xindex> to tensor<1xindex> |
| %2 = "mhlo.dynamic_broadcast_in_dim"(%arg0, %1) { broadcast_dimensions = dense<0> : tensor<1xi64> } : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32> |
| // CHECK: %[[RES:.*]] = tensor.cast %[[ARG]] : tensor<*xf32> to tensor<?xf32> |
| // CHECK: return %[[RES]] : tensor<?xf32> |
| return %2 : tensor<?xf32> |
| } |
| |
| // CHECK-LABEL: func @dynamic_broadcast_in_dim_to_same_shape_4 |
| func @dynamic_broadcast_in_dim_to_same_shape_4(%arg0: tensor<*xf32>) -> tensor<?xf32> { |
| // CHECK-SAME: %[[ARG:.*]]: tensor<*xf32> |
| %0 = shape.shape_of %arg0 : tensor<*xf32> -> !shape.shape |
| %1 = shape.to_extent_tensor %0 : !shape.shape -> tensor<?xindex> |
| %2 = tensor.cast %1 : tensor<?xindex> to tensor<1xindex> |
| %3 = "mhlo.dynamic_broadcast_in_dim"(%arg0, %2) { broadcast_dimensions = dense<0> : tensor<1xi64> } : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32> |
| // CHECK: %[[RES:.*]] = tensor.cast %[[ARG]] : tensor<*xf32> to tensor<?xf32> |
| // CHECK: return %[[RES]] : tensor<?xf32> |
| return %3 : tensor<?xf32> |
| } |
| |
| // CHECK-LABEL: func @broadcast_in_dim_constant_fold_0d |
| func @broadcast_in_dim_constant_fold_0d() -> tensor<1x64x224x224xf32> { |
| %cst = mhlo.constant dense<0.000000e+00> : tensor<f32> |
| %b = "mhlo.broadcast_in_dim"(%cst) {broadcast_dimensions = dense<[]> : tensor<0xi64>} : (tensor<f32>) -> tensor<1x64x224x224xf32> |
| return %b : tensor<1x64x224x224xf32> |
| } |
| // CHECK-NEXT: %[[CST:.*]] = mhlo.constant dense<0.000000e+00> : tensor<1x64x224x224xf32> |
| // CHECK-NEXT: return %[[CST]] : tensor<1x64x224x224xf32> |
| |
| // CHECK-LABEL: func @broadcast_in_dim_constant_fold |
| func @broadcast_in_dim_constant_fold() -> tensor<1x64x4x4xf32> { |
| %cst = mhlo.constant dense<0.000000e+00> : tensor<4x4xf32> |
| %b = "mhlo.broadcast_in_dim"(%cst) {broadcast_dimensions = dense<[2, 3]> : tensor<2xi64>} : (tensor<4x4xf32>) -> tensor<1x64x4x4xf32> |
| return %b : tensor<1x64x4x4xf32> |
| } |
| // CHECK-NEXT: %[[CST:.*]] = mhlo.constant dense<0.000000e+00> : tensor<1x64x4x4xf32> |
| // CHECK-NEXT: return %[[CST]] : tensor<1x64x4x4xf32> |
| |
| // CHECK-LABEL: @complex_expand_fold |
| func @complex_expand_fold(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { |
| %0 = "mhlo.complex"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> (tensor<4xcomplex<f32>>) |
| %1 = "mhlo.real"(%0) : (tensor<4xcomplex<f32>>) -> (tensor<4xf32>) |
| %2 = "mhlo.imag"(%0) : (tensor<4xcomplex<f32>>) -> (tensor<4xf32>) |
| // CHECK: return %arg0, %arg1 |
| return %1, %2 : tensor<4xf32>, tensor<4xf32> |
| } |
| |
| // CHECK-LABEL: @complex_collapse_fold |
| func @complex_collapse_fold(%arg0: tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>> { |
| %0 = "mhlo.real"(%arg0) : (tensor<4xcomplex<f32>>) -> (tensor<4xf32>) |
| %1 = "mhlo.imag"(%arg0) : (tensor<4xcomplex<f32>>) -> (tensor<4xf32>) |
| %2 = "mhlo.complex"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xcomplex<f32>> |
| // CHECK: return %arg0 |
| return %2 : tensor<4xcomplex<f32>> |
| } |
| |
| // CHECK-LABEL: @dynamic_iota_is_static |
| func @dynamic_iota_is_static(%arg0 : tensor<1xindex>) -> tensor<4xi32> { |
| // CHECK: [[RESULT:%.*]] = "mhlo.iota" |
| // CHECK: return [[RESULT]] |
| %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 0 : i64} : (tensor<1xindex>) -> tensor<4xi32> |
| return %0 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: @dynamic_iota_broadcast |
| func @dynamic_iota_broadcast(%arg0 : tensor<2xindex>) -> tensor<5x?xi32> { |
| // CHECK: [[IOTA:%.+]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<5xi32> |
| // CHECK: [[BROADCAST:%.+]] = "mhlo.dynamic_broadcast_in_dim"([[IOTA]], %arg0) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>, tensor<2xindex>) -> tensor<5x?xi32> |
| %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 0 : i64} : (tensor<2xindex>) -> tensor<5x?xi32> |
| |
| // CHECK: return [[BROADCAST]] |
| return %0 : tensor<5x?xi32> |
| } |
| |
| // CHECK-LABEL: @dynamic_iota_broadcast_second |
| func @dynamic_iota_broadcast_second(%arg0 : tensor<2xindex>) -> tensor<5x?xi32> { |
| // CHECK-NEXT: [[CAST1:%.+]] = index_cast %arg0 : tensor<2xindex> to tensor<2xi64> |
| // CHECK-NEXT: [[SLICE:%.+]] = "mhlo.slice"([[CAST1]]) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xi64>) -> tensor<1xi64> |
| // CHECK-NEXT: [[CAST2:%.+]] = index_cast [[SLICE]] : tensor<1xi64> to tensor<1xindex> |
| // CHECK-NEXT: [[IOTA:%.+]] = "mhlo.dynamic_iota"([[CAST2]]) {iota_dimension = 0 : i64} : (tensor<1xindex>) -> tensor<?xi32> |
| // CHECK-NEXT: [[BROADCAST:%.+]] = "mhlo.dynamic_broadcast_in_dim"([[IOTA]], %arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<?xi32>, tensor<2xindex>) -> tensor<5x?xi32> |
| %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 1 : i64} : (tensor<2xindex>) -> tensor<5x?xi32> |
| |
| // CHECK: return [[BROADCAST]] |
| return %0 : tensor<5x?xi32> |
| } |
| |
| // CHECK-LABEL: @dynamic_iota_constant |
| func @dynamic_iota_constant(%arg0 : tensor<2xindex>) -> tensor<1x?xi32> { |
| // CHECK: [[IOTA:%.+]] = mhlo.constant dense<0> : tensor<1xi32> |
| // CHECK: [[BROADCAST:%.+]] = "mhlo.dynamic_broadcast_in_dim"([[IOTA]], %arg0) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1xi32>, tensor<2xindex>) -> tensor<1x?xi32> |
| %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 0 : i64} : (tensor<2xindex>) -> tensor<1x?xi32> |
| |
| // CHECK: return [[BROADCAST]] |
| return %0 : tensor<1x?xi32> |
| } |
| |
| // CHECK-LABEL: @iota_constant |
| func @iota_constant() -> tensor<1xi32> { |
| // CHECK: [[CONST:%.+]] = mhlo.constant dense<0> : tensor<1xi32> |
| %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<1xi32> |
| |
| // CHECK: return [[CONST]] : tensor<1xi32> |
| return %0 : tensor<1xi32> |
| } |
| |
| // CHECK-LABEL: @iota_constant_multi |
| func @iota_constant_multi() -> tensor<1x4xi32> { |
| // CHECK: [[CONST:%.+]] = mhlo.constant dense<0> : tensor<1x4xi32> |
| %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<1x4xi32> |
| |
| // CHECK: return [[CONST]] : tensor<1x4xi32> |
| return %0 : tensor<1x4xi32> |
| } |
| |
| // CHECK-LABEL: @iota_not_lowered_to_constant |
| func @iota_not_lowered_to_constant() -> tensor<4xi32> { |
| // CHECK: [[RESULT:%.*]] = "mhlo.iota" |
| // CHECK: return [[RESULT]] |
| %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xi32> |
| return %0 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: @iota_broadcast |
| func @iota_broadcast() -> tensor<5x4xi32> { |
| // CHECK: [[IOTA:%.+]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<5xi32> |
| // CHECK: [[RESULT:%.+]] = "mhlo.broadcast_in_dim"([[IOTA]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>) -> tensor<5x4xi32> |
| %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<5x4xi32> |
| |
| return %0 : tensor<5x4xi32> |
| } |
| |
| // CHECK-LABEL: @iota_broadcast |
| func @iota_broadcast_second() -> tensor<5x4xi32> { |
| // CHECK: [[IOTA:%.+]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xi32> |
| // CHECK: [[RESULT:%.+]] = "mhlo.broadcast_in_dim"([[IOTA]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<4xi32>) -> tensor<5x4xi32> |
| %0 = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<5x4xi32> |
| |
| return %0 : tensor<5x4xi32> |
| } |
| |
| // CHECK-LABEL: @unary_einsum |
| func @unary_einsum(%arg0: tensor<2x3xf32>) -> tensor<2x2xf32> { |
| // CHECK: %[[ONE:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32> |
| // CHECK: "mhlo.einsum"(%[[ONE]], %arg0) {einsum_config = ",ab->aa"} |
| %0 = "mhlo.unary_einsum"(%arg0) {einsum_config = "ab->aa"} : (tensor<2x3xf32>) -> tensor<2x2xf32> |
| return %0 : tensor<2x2xf32> |
| } |
| |
| // CHECK-LABEL: func @fold_copy |
| // CHECK-SAME: [[ARG:%[a-zA-Z0-9]+]] |
| func @fold_copy(%arg : tensor<1x4xf32>) -> tensor<1x4xf32> { |
| // CHECK: return [[ARG]] |
| %0 = "mhlo.copy"(%arg) : (tensor<1x4xf32>) -> tensor<1x4xf32> |
| return %0 : tensor<1x4xf32> |
| } |
| |
| // CHECK-LABEL: func @dynamic_reshape_not_actually_dynamic |
| func @dynamic_reshape_not_actually_dynamic(%arg0: tensor<4xf32>, %shape: tensor<2xindex>) -> tensor<4x1xf32> { |
| // CHECK: mhlo.reshape |
| %0 = "mhlo.dynamic_reshape"(%arg0, %shape) : (tensor<4xf32>, tensor<2xindex>) -> tensor<4x1xf32> |
| return %0 : tensor<4x1xf32> |
| } |
| |
| // CHECK-LABEL: func @shape_of_dynamic_reshape |
| // CHECK-SAME: [[ARG0:%[a-zA-Z0-9]+]] |
| // CHECK-SAME: [[ARG1:%[a-zA-Z0-9]+]] |
| func @shape_of_dynamic_reshape(%arg0: tensor<*xf32>, %shape: tensor<2xindex>) -> tensor<2xindex> { |
| // CHECK: return [[ARG1]] |
| %0 = "mhlo.dynamic_reshape"(%arg0, %shape) : (tensor<*xf32>, tensor<2xindex>) -> tensor<?x?xf32> |
| %1 = shape.shape_of %0 : tensor<?x?xf32> -> tensor<2xindex> |
| return %1 : tensor<2xindex> |
| } |
| |
| // CHECK-LABEL: func @dynamic_reshape_rank_1_to_rank_1 |
| // CHECK-SAME: [[ARG0:%[a-zA-Z0-9]+]] |
| func @dynamic_reshape_rank_1_to_rank_1(%arg0: tensor<?xcomplex<f32>>, |
| %shape: tensor<?xindex>) -> tensor<?xf32> { |
| // CHECK: [[RES:%[a-zA-Z0-9]+]] = "mhlo.real"([[ARG0]]) : (tensor<?xcomplex<f32>>) -> tensor<?xf32> |
| // CHECK: return [[RES]] |
| %0 = "mhlo.real"(%arg0): (tensor<?xcomplex<f32>>) -> tensor<?xf32> |
| %1 = shape.shape_of %arg0 : tensor<?xcomplex<f32>> -> tensor<1xindex> |
| %2 = shape.num_elements %1 : tensor<1xindex> -> index |
| %3 = tensor.from_elements %2 : tensor<1xindex> |
| %4 = "mhlo.dynamic_reshape"(%0, %3) |
| : (tensor<?xf32>, tensor<1xindex>) -> tensor<?xf32> |
| return %4 : tensor<?xf32> |
| } |
| |
| // CHECK-LABEL: func @dynamic_reshape_of_dynamic_reshape |
| // CHECK-SAME: [[ARG0:%[a-zA-Z0-9]+]] |
| // CHECK-SAME: [[ARG1:%[a-zA-Z0-9]+]] |
| func @dynamic_reshape_of_dynamic_reshape(%arg0: tensor<?xf16>, %shape: tensor<?xindex>) -> tensor<?xf16> { |
| // CHECK: return [[ARG0]] |
| %0 = "mhlo.dynamic_reshape"(%arg0, %shape) : (tensor<?xf16>, tensor<?xindex>) -> tensor<*xf16> |
| %1 = shape.shape_of %0 : tensor<*xf16> -> tensor<?xindex> |
| %2 = shape.num_elements %1 : tensor<?xindex> -> index |
| %3 = tensor.from_elements %2 : tensor<1xindex> |
| %4 = "mhlo.dynamic_reshape"(%0, %3) : (tensor<*xf16>, tensor<1xindex>) -> tensor<?xf16> |
| return %4 : tensor<?xf16> |
| } |
| |
| // CHECK-LABEL: do_not_dce_while_with_outfeed |
| func @do_not_dce_while_with_outfeed(%arg0: tensor<i64>) -> tensor<i64> { |
| // CHECK: mhlo.while |
| %0 = "mhlo.while"(%arg0) ( { |
| ^bb0(%arg1: tensor<i64>): |
| %1 = "mhlo.compare"(%arg1, %arg1) {comparison_direction = "LT"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| "mhlo.return"(%1) : (tensor<i1>) -> () |
| }, { |
| ^bb0(%arg1: tensor<i64>): |
| %1 = "mhlo.create_token"() : () -> !mhlo.token |
| // Side-effecting op outfeed present inside while. |
| %2 = "mhlo.outfeed"(%arg1, %1) {outfeed_config = ""} : (tensor<i64>, !mhlo.token) -> !mhlo.token |
| "mhlo.return"(%arg1) : (tensor<i64>) -> () |
| }) : (tensor<i64>) -> tensor<i64> |
| |
| return %arg0 : tensor<i64> |
| } |
| |
| // CHECK-LABEL: dce_while_without_side_effect |
| func @dce_while_without_side_effect(%arg0: tensor<i64>) -> tensor<i64> { |
| // CHECK-NOT: mhlo.while |
| %0 = "mhlo.while"(%arg0) ( { |
| ^bb0(%arg1: tensor<i64>): |
| %1 = "mhlo.compare"(%arg1, %arg1) {comparison_direction = "LT"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| "mhlo.return"(%1) : (tensor<i1>) -> () |
| }, { |
| ^bb0(%arg1: tensor<i64>): |
| %1 = "mhlo.create_token"() : () -> !mhlo.token |
| "mhlo.return"(%arg1) : (tensor<i64>) -> () |
| }) : (tensor<i64>) -> tensor<i64> |
| |
| return %arg0 : tensor<i64> |
| } |
| |
| // CHECK-LABEL: fold_compare_same_eq |
| func @fold_compare_same_eq(%arg0: tensor<i64>) -> tensor<i1> { |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "EQ"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| return %0 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_same_le |
| func @fold_compare_same_le(%arg0: tensor<i64>) -> tensor<i1> { |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "LE"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| return %0 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_same_ge |
| func @fold_compare_same_ge(%arg0: tensor<i64>) -> tensor<i1> { |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "GE"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| return %0 : tensor<i1> |
| } |
| // CHECK-LABEL: fold_compare_same_ne |
| func @fold_compare_same_ne(%arg0: tensor<i64>) -> tensor<i1> { |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "NE"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| return %0 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_same_lt |
| func @fold_compare_same_lt(%arg0: tensor<i64>) -> tensor<i1> { |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "LT"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| return %0 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_same_gt |
| func @fold_compare_same_gt(%arg0: tensor<i64>) -> tensor<i1> { |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "GT"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| return %0 : tensor<i1> |
| } |
| |
| // Address NaN != NaN. |
| // CHECK-LABEL: dont_fold_compare_same_eq_float |
| func @dont_fold_compare_same_eq_float(%arg0: tensor<f16>) -> tensor<i1> { |
| // CHECK: %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "EQ"} : (tensor<f16>, tensor<f16>) -> tensor<i1> |
| %0 = "mhlo.compare"(%arg0, %arg0) {comparison_direction = "EQ"} : (tensor<f16>, tensor<f16>) -> tensor<i1> |
| return %0 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_eq |
| func @fold_compare_false_eq() -> tensor<i1> { |
| %0 = mhlo.constant dense<0> : tensor<i32> |
| %1 = mhlo.constant dense<1> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| // CHECK-LABEL: fold_compare_true_eq |
| func @fold_compare_true_eq() -> tensor<i1> { |
| %0 = mhlo.constant dense<1> : tensor<i32> |
| %1 = mhlo.constant dense<1> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_eq_float |
| func @fold_compare_false_eq_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<0.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "EQ"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_eq_float |
| func @fold_compare_true_eq_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<1.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "EQ"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_ne |
| func @fold_compare_false_ne() -> tensor<i1> { |
| %0 = mhlo.constant dense<1> : tensor<i32> |
| %1 = mhlo.constant dense<1> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "NE"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_ne |
| func @fold_compare_true_ne() -> tensor<i1> { |
| %0 = mhlo.constant dense<1> : tensor<i32> |
| %1 = mhlo.constant dense<0> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "NE"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_ne_float |
| func @fold_compare_false_ne_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<1.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "NE"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_ne_float |
| func @fold_compare_true_ne_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<0.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "NE"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_lt |
| func @fold_compare_false_lt() -> tensor<i1> { |
| %0 = mhlo.constant dense<1> : tensor<i32> |
| %1 = mhlo.constant dense<1> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_lt |
| func @fold_compare_true_lt() -> tensor<i1> { |
| %0 = mhlo.constant dense<0> : tensor<i32> |
| %1 = mhlo.constant dense<1> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_lt_float |
| func @fold_compare_false_lt_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<1.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LT"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_lt_float |
| func @fold_compare_true_lt_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<0.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LT"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_le |
| func @fold_compare_false_le() -> tensor<i1> { |
| %0 = mhlo.constant dense<1> : tensor<i32> |
| %1 = mhlo.constant dense<0> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LE"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_le |
| func @fold_compare_true_le() -> tensor<i1> { |
| %0 = mhlo.constant dense<1> : tensor<i32> |
| %1 = mhlo.constant dense<1> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LE"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_le_float |
| func @fold_compare_false_le_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<1.> : tensor<f32> |
| %1 = mhlo.constant dense<0.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LE"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_le_float |
| func @fold_compare_true_le_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<1.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "LE"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_gt |
| func @fold_compare_false_gt() -> tensor<i1> { |
| %0 = mhlo.constant dense<0> : tensor<i32> |
| %1 = mhlo.constant dense<0> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_gt |
| func @fold_compare_true_gt() -> tensor<i1> { |
| %0 = mhlo.constant dense<1> : tensor<i32> |
| %1 = mhlo.constant dense<0> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_gt_float |
| func @fold_compare_false_gt_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<0.> : tensor<f32> |
| %1 = mhlo.constant dense<0.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_gt_float |
| func @fold_compare_true_gt_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<1.> : tensor<f32> |
| %1 = mhlo.constant dense<0.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_ge |
| func @fold_compare_false_ge() -> tensor<i1> { |
| %0 = mhlo.constant dense<0> : tensor<i32> |
| %1 = mhlo.constant dense<1> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GE"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_ge |
| func @fold_compare_true_ge() -> tensor<i1> { |
| %0 = mhlo.constant dense<0> : tensor<i32> |
| %1 = mhlo.constant dense<0> : tensor<i32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GE"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_false_ge_float |
| func @fold_compare_false_ge_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<0.> : tensor<f32> |
| %1 = mhlo.constant dense<1.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<false> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GE"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: fold_compare_true_ge_float |
| func @fold_compare_true_ge_float() -> tensor<i1> { |
| %0 = mhlo.constant dense<0.> : tensor<f32> |
| %1 = mhlo.constant dense<0.> : tensor<f32> |
| // CHECK: %0 = mhlo.constant dense<true> : tensor<i1> |
| %2 = "mhlo.compare"(%0, %1) {comparison_direction = "GE"} : (tensor<f32>, tensor<f32>) -> tensor<i1> |
| return %2 : tensor<i1> |
| } |
| |
| // CHECK-LABEL: unpack_repack_same_tuple |
| // CHECK-SAME: ([[ARG0:%.*]]: tuple<tensor<i32>, !mhlo.token, tensor<f32>>) |
| func @unpack_repack_same_tuple(%arg0: tuple<tensor<i32>, !mhlo.token, tensor<f32>>) -> tuple<tensor<i32>, !mhlo.token, tensor<f32>> { |
| %0 = "mhlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i32>, !mhlo.token, tensor<f32>>) -> tensor<i32> |
| %1 = "mhlo.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tensor<i32>, !mhlo.token, tensor<f32>>) -> !mhlo.token |
| %2 = "mhlo.get_tuple_element"(%arg0) {index = 2 : i32} : (tuple<tensor<i32>, !mhlo.token, tensor<f32>>) -> tensor<f32> |
| %3 = "mhlo.tuple"(%0, %1, %2) : (tensor<i32>, !mhlo.token, tensor<f32>) -> tuple<tensor<i32>, !mhlo.token, tensor<f32>> |
| |
| // CHECK: return [[ARG0]] |
| return %3 : tuple<tensor<i32>, !mhlo.token, tensor<f32>> |
| } |
| |
| // CHECK-LABEL: unpack_repack_same_tuple_single_element |
| // CHECK-SAME: ([[ARG0:%.*]]: tuple<tensor<i32>>) |
| func @unpack_repack_same_tuple_single_element(%arg0: tuple<tensor<i32>>) -> tuple<tensor<i32>> { |
| %0 = "mhlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i32>>) -> tensor<i32> |
| %3 = "mhlo.tuple"(%0) : (tensor<i32>) -> tuple<tensor<i32>> |
| |
| // CHECK: return [[ARG0]] |
| return %3 : tuple<tensor<i32>> |
| } |
| |
| // CHECK-LABEL: func @erase_dead_lhlo_constant |
| func @erase_dead_lhlo_constant() { |
| %M = memref.alloc() : memref<256x1024xf32> |
| // CHECK-NEXT: return |
| "lmhlo.constant"(%M) {value = dense<0.0> : tensor<f32>} : (memref<256x1024xf32>) -> () |
| memref.dealloc %M : memref<256x1024xf32> |
| return |
| } |
| |
| // A negative test for dead lhlo constant op erasure. |
| // CHECK-LABEL: func @erase_dead_lhlo_constant_negative |
| func @erase_dead_lhlo_constant_negative(%M : memref<4xf32>) -> memref<256x1024xf32> { |
| // CHECK-NEXT: lmhlo.constant |
| "lmhlo.constant"(%M) {value = dense<0.0> : tensor<f32>} : (memref<4xf32>) -> () |
| // CHECK-NEXT: memref.alloc |
| // CHECK-NEXT: lmhlo.constant |
| %N = memref.alloc() : memref<256x1024xf32> |
| "lmhlo.constant"(%N) {value = dense<0.0> : tensor<f32>} : (memref<256x1024xf32>) -> () |
| return %N : memref<256x1024xf32> |
| } |
| |
| // CHECK-LABEL: func @fold_get_dimension_size |
| func @fold_get_dimension_size(%I: tensor<1x128x512xf32>) -> tensor<i32> { |
| %size = "mhlo.get_dimension_size"(%I) {dimension = 2 : i64} : (tensor<1x128x512xf32>) -> tensor<i32> |
| return %size : tensor<i32> |
| // CHECK-NEXT: %[[C:.*]] = mhlo.constant dense<512> : tensor<i32> |
| // CHECK-NEXT: return %[[C]] |
| } |
| |
| // CHECK-LABEL: func @fold_set_dimension_size |
| // CHECK-SAME: (%[[I:.*]]: tensor<1x128x512xf32>) |
| func @fold_set_dimension_size(%I: tensor<1x128x512xf32>) -> tensor<1x128x512xf32> { |
| %dim = mhlo.constant dense<512> : tensor<i32> |
| %result = "mhlo.set_dimension_size"(%I, %dim) {dimension = 2 : i64} : (tensor<1x128x512xf32>, tensor<i32>) -> tensor<1x128x512xf32> |
| return %result : tensor<1x128x512xf32> |
| |
| // CHECK-NEXT: return %[[I]] |
| } |
| |
| // CHECK-LABEL: func @fold_select_same |
| func @fold_select_same(%arg0 : tensor<f32>, %arg1 : tensor<i1>) -> tensor<f32> { |
| %1 = "mhlo.select"(%arg1, %arg0, %arg0) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| // CHECK: return %arg0 |
| return %1 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: func @fold_select_first |
| func @fold_select_first(%arg0 : tensor<f32>, %arg1 : tensor<f32>) -> tensor<f32> { |
| %0 = mhlo.constant dense<1> : tensor<i1> |
| %1 = "mhlo.select"(%0, %arg0, %arg1) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| // CHECK: return %arg0 |
| return %1 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: func @fold_select_second |
| func @fold_select_second(%arg0 : tensor<f32>, %arg1 : tensor<f32>) -> tensor<f32> { |
| %0 = mhlo.constant dense<0> : tensor<i1> |
| %1 = "mhlo.select"(%0, %arg0, %arg1) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| // CHECK: return %arg1 |
| return %1 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: func @fold_select_vector |
| func @fold_select_vector(%arg0 : tensor<4xf32>, %arg1 : tensor<4xf32>) -> tensor<4xf32> { |
| %0 = mhlo.constant dense<1> : tensor<4xi1> |
| %1 = "mhlo.select"(%0, %arg0, %arg1) : (tensor<4xi1>, tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> |
| // CHECK: return %arg0 |
| return %1 : tensor<4xf32> |
| } |
| |
| // CHECK-LABEL: gather_to_slice |
| func @gather_to_slice(%arg0: tensor<5x6x7xf32>) -> tensor<3x6x5xf32> { |
| %0 = constant dense<[1, 2]> : tensor<2xi32> |
| %1 = "mhlo.gather"(%arg0, %0) { |
| dimension_numbers = {collapsed_slice_dims = dense<> : tensor<0xi64>, |
| index_vector_dim = 0 : i64, |
| offset_dims = dense<[0, 1, 2]> : tensor<3xi64>, |
| start_index_map = dense<[0, 2]> : tensor<2xi64>}, |
| indices_are_sorted = false, |
| slice_sizes = dense<[3, 6, 5]> : tensor<3xi64>} : (tensor<5x6x7xf32>, tensor<2xi32>) -> tensor<3x6x5xf32> |
| return %1 : tensor<3x6x5xf32> |
| // CHECK: %[[RET:.*]] = "mhlo.slice"(%arg0) {limit_indices = dense<[4, 6, 7]> : tensor<3xi64>, start_indices = dense<[1, 0, 2]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<5x6x7xf32>) -> tensor<3x6x5xf32> |
| // CHECK: return %[[RET]] : tensor<3x6x5xf32> |
| } |
| |
| // CHECK-LABEL: gather_scalar_index_to_slice |
| func @gather_scalar_index_to_slice(%arg0: tensor<5x6x7xf32>) -> tensor<5x6x4xf32> { |
| %0 = constant dense<1> : tensor<i32> |
| %1 = "mhlo.gather"(%arg0, %0) { |
| dimension_numbers = {collapsed_slice_dims = dense<> : tensor<0xi64>, |
| index_vector_dim = 0 : i64, |
| offset_dims = dense<[0, 1, 2]> : tensor<3xi64>, |
| start_index_map = dense<[2]> : tensor<1xi64>}, |
| indices_are_sorted = false, |
| slice_sizes = dense<[5, 6, 4]> : tensor<3xi64>} : (tensor<5x6x7xf32>, tensor<i32>) -> tensor<5x6x4xf32> |
| return %1 : tensor<5x6x4xf32> |
| // CHECK: %[[RET:.*]] = "mhlo.slice"(%arg0) {limit_indices = dense<[5, 6, 5]> : tensor<3xi64>, start_indices = dense<[0, 0, 1]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<5x6x7xf32>) -> tensor<5x6x4xf32> |
| // CHECK: return %[[RET]] : tensor<5x6x4xf32> |
| } |
| |
| // CHECK-LABEL: gather_to_slice_reshape |
| func @gather_to_slice_reshape(%arg0: tensor<5x6x7xf32>) -> tensor<3x6xf32> { |
| %0 = constant dense<[1, 2]> : tensor<2xi32> |
| %1 = "mhlo.gather"(%arg0, %0) { |
| dimension_numbers = {collapsed_slice_dims = dense<[2]> : tensor<1xi64>, |
| index_vector_dim = 0 : i64, |
| offset_dims = dense<[0, 1, 2]> : tensor<3xi64>, |
| start_index_map = dense<[0, 2]> : tensor<2xi64>}, |
| indices_are_sorted = false, |
| slice_sizes = dense<[3, 6, 1]> : tensor<3xi64>} : (tensor<5x6x7xf32>, tensor<2xi32>) -> tensor<3x6xf32> |
| return %1 : tensor<3x6xf32> |
| // CHECK: %[[V0:.*]] = "mhlo.slice"(%arg0) {limit_indices = dense<[4, 6, 3]> : tensor<3xi64>, start_indices = dense<[1, 0, 2]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<5x6x7xf32>) -> tensor<3x6x1xf32> |
| // CHECK: %[[V1:.*]] = "mhlo.reshape"(%[[V0]]) : (tensor<3x6x1xf32>) -> tensor<3x6xf32> |
| // CHECK: return %[[V1]] : tensor<3x6xf32> |
| } |
| |
| // CHECK-LABEL: func @fold_and_same |
| func @fold_and_same(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = "mhlo.and"(%arg0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %arg0 |
| return %0 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_and_ones |
| func @fold_and_ones(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<-1> : tensor<4xi32> |
| %1 = "mhlo.and"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %arg0 |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_and_zeros |
| func @fold_and_zeros(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<0> : tensor<4xi32> |
| %1 = "mhlo.and"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %0 |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_and_constant |
| func @fold_and_constant(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<7> : tensor<4xi32> |
| // CHECK: mhlo.and |
| %1 = "mhlo.and"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_and_constants |
| func @fold_and_constants() -> tensor<4xi32> { |
| %0 = mhlo.constant dense<[0, 1, 6, 3]> : tensor<4xi32> |
| %1 = mhlo.constant dense<[7, 3, 7, 2]> : tensor<4xi32> |
| %2 = "mhlo.and"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: %0 = mhlo.constant dense<[0, 1, 6, 2]> : tensor<4xi32> |
| // CHECK: return %0 |
| return %2 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_or_same |
| func @fold_or_same(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = "mhlo.or"(%arg0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %arg0 |
| return %0 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_or_ones |
| func @fold_or_ones(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<-1> : tensor<4xi32> |
| %1 = "mhlo.or"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %0 |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_or_zeros |
| func @fold_or_zeros(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<0> : tensor<4xi32> |
| %1 = "mhlo.or"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %arg0 |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_or_constant |
| func @fold_or_constant(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<7> : tensor<4xi32> |
| // CHECK: mhlo.or |
| %1 = "mhlo.or"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_or_zeros_right |
| func @fold_or_zeros_right(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<0> : tensor<4xi32> |
| %1 = "mhlo.or"(%arg0, %0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %arg0 |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_or_zeros_constants |
| func @fold_or_zeros_constants() -> tensor<4xi32> { |
| %0 = mhlo.constant dense<[0, 1, 6, 3]> : tensor<4xi32> |
| %1 = mhlo.constant dense<[7, 3, 7, 2]> : tensor<4xi32> |
| %2 = "mhlo.or"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: %0 = mhlo.constant dense<[7, 3, 7, 3]> : tensor<4xi32> |
| // CHECK: return %0 |
| return %2 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_xor_same |
| func @fold_xor_same(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = "mhlo.xor"(%arg0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: %0 = mhlo.constant dense<0> : tensor<4xi32> |
| // CHECK: return %0 |
| return %0 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_xor_ones_left |
| func @fold_xor_ones_left(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<-1> : tensor<4xi32> |
| // CHECK: mhlo.xor |
| %1 = "mhlo.xor"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_xor_ones_right |
| func @fold_xor_ones_right(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<-1> : tensor<4xi32> |
| // CHECK: mhlo.xor |
| %1 = "mhlo.xor"(%arg0, %0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_xor_zeros_left |
| func @fold_xor_zeros_left(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<0> : tensor<4xi32> |
| %1 = "mhlo.xor"(%0, %arg0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %arg0 |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_xor_zeros_right |
| func @fold_xor_zeros_right(%arg0 : tensor<4xi32>) -> tensor<4xi32> { |
| %0 = mhlo.constant dense<0> : tensor<4xi32> |
| %1 = "mhlo.xor"(%arg0, %0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: return %arg0 |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_xor_zeros_constants |
| func @fold_xor_zeros_constants() -> tensor<4xi32> { |
| %0 = mhlo.constant dense<[0, 1, 6, 3]> : tensor<4xi32> |
| %1 = mhlo.constant dense<[7, 3, 7, 2]> : tensor<4xi32> |
| %2 = "mhlo.xor"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> |
| // CHECK: %0 = mhlo.constant dense<[7, 2, 1, 1]> : tensor<4xi32> |
| // CHECK: return %0 |
| return %2 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_negate_int |
| func @fold_negate_int() -> tensor<4xi32> { |
| %0 = mhlo.constant dense<[0, 1, 6, -3]> : tensor<4xi32> |
| // CHECK: mhlo.constant dense<[0, -1, -6, 3]> |
| %1 = "mhlo.negate"(%0) : (tensor<4xi32>) -> tensor<4xi32> |
| return %1 : tensor<4xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_negate_float |
| func @fold_negate_float() -> tensor<4xf32> { |
| %0 = mhlo.constant dense<[0., 1., 6., -3.]> : tensor<4xf32> |
| // CHECK: mhlo.constant dense<[-0.000000e+00, -1.000000e+00, -6.000000e+00, 3.000000e+00]> |
| %1 = "mhlo.negate"(%0) : (tensor<4xf32>) -> tensor<4xf32> |
| return %1 : tensor<4xf32> |
| } |
| |
| // CHECK-LABEL func @fold_not() |
| func @fold_not() -> tensor<2x2xi1> { |
| %0 = mhlo.constant dense<[[true, false], [true, false]]> : tensor<2x2xi1> |
| // CHECK{LITERAL}: mhlo.constant dense<[[false, true], [false, true]]> : tensor<2x2xi1> |
| %1 = "mhlo.not"(%0) : (tensor<2x2xi1>) -> tensor<2x2xi1> |
| return %1 : tensor<2x2xi1> |
| } |
| |
| // CHECK-LABEL func @fold_not_i32() |
| func @fold_not_i32() -> tensor<2x2xi32> { |
| %0 = mhlo.constant dense<[[42, -12], [1, 0]]> : tensor<2x2xi32> |
| // CHECK-LITERAL: mhlo.constant dense<[[0, 0], [0, 1]]> : tensor<2x2xi32> |
| %1 = "mhlo.not"(%0) : (tensor<2x2xi32>) -> tensor<2x2xi32> |
| return %1 : tensor<2x2xi32> |
| } |
| |
| // CHECK-LABEL: func @fold_sqrt_f32_constants |
| func @fold_sqrt_f32_constants() -> tensor<4xf32> { |
| %0 = mhlo.constant dense<1.0> : tensor<4xf32> |
| %1 = "mhlo.sqrt"(%0) : (tensor<4xf32>) -> tensor<4xf32> |
| // CHECK: mhlo.constant dense<1.000000e+00> : tensor<4xf32> |
| // CHECK-NOT: mhlo.sqrt |
| return %1 : tensor<4xf32> |
| } |
| |
| // CHECK-LABEL: func @fold_sqrt_f64_constants |
| func @fold_sqrt_f64_constants() -> tensor<4xf64> { |
| %0 = mhlo.constant dense<[1.0, 4.0, 9.0, 16.0]> : tensor<4xf64> |
| %1 = "mhlo.sqrt"(%0) : (tensor<4xf64>) -> tensor<4xf64> |
| // CHECK: mhlo.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : tensor<4xf64> |
| // CHECK-NOT: mhlo.sqrt |
| return %1 : tensor<4xf64> |
| } |
| |
| // CHECK-LABEL: func @not_fold_sqrt_neg_constants |
| func @not_fold_sqrt_neg_constants() -> tensor<4xf32> { |
| %0 = mhlo.constant dense<-1.0> : tensor<4xf32> |
| %1 = "mhlo.sqrt"(%0) : (tensor<4xf32>) -> tensor<4xf32> |
| // CHECK: mhlo.constant dense<-1.000000e+00> : tensor<4xf32> |
| // CHECK: mhlo.sqrt |
| return %1 : tensor<4xf32> |
| } |
| |
| // CHECK-LABEL: func @fold_if_true( |
| // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: ) |
| func @fold_if_true(%arg0 : tensor<f32>, %arg1 : tensor<f32>) -> tensor<f32> { |
| // CHECK-NOT: mhlo.if |
| // CHECK: return %[[ARG0]] |
| %true = mhlo.constant dense<true> : tensor<i1> |
| %0 = "mhlo.if"(%true, %arg0, %arg1) ( { |
| ^bb0(%bbarg0: tensor<f32>): |
| "mhlo.return"(%bbarg0) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg1: tensor<f32>): |
| "mhlo.return"(%bbarg1) : (tensor<f32>) -> () |
| }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| return %0 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: func @fold_if_false( |
| // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: ) |
| func @fold_if_false(%arg0 : tensor<f32>, %arg1 : tensor<f32>) -> tensor<f32> { |
| // CHECK-NOT: mhlo.if |
| // CHECK: return %[[ARG1]] |
| %false = mhlo.constant dense<false> : tensor<i1> |
| %0 = "mhlo.if"(%false, %arg0, %arg1) ( { |
| ^bb0(%bbarg0: tensor<f32>): |
| "mhlo.return"(%bbarg0) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg1: tensor<f32>): |
| "mhlo.return"(%bbarg1) : (tensor<f32>) -> () |
| }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| return %0 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: func @fold_case( |
| // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: ) |
| func @fold_case(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<f32>) -> tensor<f32> { |
| // CHECK-NOT: mhlo.case |
| // CHECK: return %[[ARG1]] |
| %c1 = mhlo.constant dense<1> : tensor<i32> |
| %0 = "mhlo.case"(%c1, %arg0, %arg1, %arg2) ( { |
| ^bb0(%bbarg0: tensor<f32>): |
| "mhlo.return"(%bbarg0) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg1: tensor<f32>): |
| "mhlo.return"(%bbarg1) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg2: tensor<f32>): |
| "mhlo.return"(%bbarg2) : (tensor<f32>) -> () |
| }) : (tensor<i32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| return %0 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: func @fold_case_negative_index( |
| // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: ) |
| func @fold_case_negative_index(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<f32>) -> tensor<f32> { |
| // CHECK-NOT: mhlo.case |
| // CHECK: return %[[ARG2]] |
| %m1000 = mhlo.constant dense<-1000> : tensor<i32> |
| %0 = "mhlo.case"(%m1000, %arg0, %arg1, %arg2) ( { |
| ^bb0(%bbarg0: tensor<f32>): |
| "mhlo.return"(%bbarg0) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg1: tensor<f32>): |
| "mhlo.return"(%bbarg1) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg2: tensor<f32>): |
| "mhlo.return"(%bbarg2) : (tensor<f32>) -> () |
| }) : (tensor<i32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| return %0 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: func @fold_case_oob_index( |
| // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]] |
| // CHECK-SAME: ) |
| func @fold_case_oob_index(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<f32>) -> tensor<f32> { |
| // CHECK-NOT: mhlo.case |
| // CHECK: return %[[ARG2]] |
| %c1000 = mhlo.constant dense<1000> : tensor<i32> |
| %0 = "mhlo.case"(%c1000, %arg0, %arg1, %arg2) ( { |
| ^bb0(%bbarg0: tensor<f32>): |
| "mhlo.return"(%bbarg0) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg1: tensor<f32>): |
| "mhlo.return"(%bbarg1) : (tensor<f32>) -> () |
| }, { |
| ^bb0(%bbarg2: tensor<f32>): |
| "mhlo.return"(%bbarg2) : (tensor<f32>) -> () |
| }) : (tensor<i32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> |
| return %0 : tensor<f32> |
| } |
| |
| // CHECK-LABEL: @tensor_flow_scatter_v1_update |
| func @tensor_flow_scatter_v1_update() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[0, 2]> : tensor<2xi32> |
| %2 = constant dense<[[10, 20, 30], [70, 80, 90]]> : tensor<2x3xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| "mhlo.return"(%arg1) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 1 : i64, |
| inserted_window_dims = dense<0> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2xi32>, tensor<2x3xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [10, 20, 30], [4, 5, 6], [70, 80, 90] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| } |
| |
| // CHECK-LABEL: @tensor_flow_scatter_v2_update |
| func @tensor_flow_scatter_v2_update() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[0, 2]> : tensor<2xi32> |
| %2 = constant dense<[[10, 30], [40, 60], [70, 90]]> : tensor<3x2xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| "mhlo.return"(%arg1) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 1 : i64, |
| inserted_window_dims = dense<1> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<1> : tensor<1xi64>, |
| update_window_dims = dense<[0]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2xi32>, tensor<3x2xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [10, 2, 30], [40, 5, 60], [70, 8, 90] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| } |
| |
| // CHECK-LABEL: @tensor_flow_scatter_add |
| func @tensor_flow_scatter_add() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[0, 2]> : tensor<2xi32> |
| %2 = constant dense<[[10, 20, 30], [70, 80, 90]]> : tensor<2x3xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| %4 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> (tensor<i32>) |
| "mhlo.return"(%4) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 1 : i64, |
| inserted_window_dims = dense<0> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2xi32>, tensor<2x3xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [11, 22, 33], [4, 5, 6], [77, 88, 99] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| } |
| |
| // CHECK-LABEL: @tensor_flow_scatter_repeated |
| func @tensor_flow_scatter_repeated() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[1, 1]> : tensor<2xi32> |
| %2 = constant dense<[[10, 20, 30], [70, 80, 90]]> : tensor<2x3xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| %4 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> (tensor<i32>) |
| "mhlo.return"(%4) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 1 : i64, |
| inserted_window_dims = dense<0> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2xi32>, tensor<2x3xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [1, 2, 3], [84, 105, 126], [7, 8, 9] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| } |
| |
| // CHECK-LABEL: @tensor_flow_scatter_multiple_batch |
| func @tensor_flow_scatter_multiple_batch() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[[0, 2], [2, 1]]> : tensor<2x2xi32> |
| %2 = constant dense<[[[10, 30], [40, 60], [70, 90]], [[5, 5], [5, 5], [5, 5]]]> : tensor<2x3x2xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| %4 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> (tensor<i32>) |
| "mhlo.return"(%4) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 2 : i64, |
| inserted_window_dims = dense<1> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<1> : tensor<1xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2x2xi32>, tensor<2x3x2xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [11, 7, 38], [44, 10, 71], [77, 13, 104] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| } |
| |
| // CHECK-LABEL: @tensor_flow_scatter_nd |
| func @tensor_flow_scatter_nd() -> tensor<3x3x2xi32> { |
| %0 = constant dense<[[[-1, 1], [-2, 2], [-3, 3]], [[-4, 4], [-5, 5], [-6, 6]], [[-7, 7], [-8, 8], [-9, 9]]]> : tensor<3x3x2xi32> |
| %1 = constant dense<[[0, 0], [1, 0]]> : tensor<2x2xi32> |
| %2 = constant dense<[[-10, 10], [-40, 40]]> : tensor<2x2xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| "mhlo.return"(%arg1) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 1 : i64, |
| inserted_window_dims = dense<[0, 1]> : tensor<2xi64>, |
| scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3x2xi32>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<3x3x2xi32> |
| return %3 : tensor<3x3x2xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [-10, 10], [-2, 2], [-3, 3] |
| // CHECK-SAME: [-40, 40], [-5, 5], [-6, 6] |
| // CHECK-SAME: [-7, 7], [-8, 8], [-9, 9] |
| // CHECK-SAME: ]> : tensor<3x3x2xi32> |
| } |
| |
| // CHECK-LABEL: @tensor_flow_scatter_nd_index_vector |
| func @tensor_flow_scatter_nd_index_vector() -> tensor<3x3x2xi32> { |
| %0 = constant dense<[[[-1, 1], [-2, 2], [-3, 3]], [[-4, 4], [-5, 5], [-6, 6]], [[-7, 7], [-8, 8], [-9, 9]]]> : tensor<3x3x2xi32> |
| %1 = constant dense<[[0, 0], [1, 0]]> : tensor<2x2xi32> |
| %2 = constant dense<[[-10, 10], [-20, 20]]> : tensor<2x2xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| "mhlo.return"(%arg1) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 0 : i64, |
| inserted_window_dims = dense<[0, 1]> : tensor<2xi64>, |
| scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3x2xi32>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<3x3x2xi32> |
| return %3 : tensor<3x3x2xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [-20, 20], [-10, 10], [-3, 3] |
| // CHECK-SAME: [-4, 4], [-5, 5], [-6, 6] |
| // CHECK-SAME: [-7, 7], [-8, 8], [-9, 9] |
| // CHECK-SAME: ]> : tensor<3x3x2xi32> |
| } |
| |
| // CHECK-LABEL: @scatter_batch_dus |
| func @scatter_batch_dus() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[[2, 1], [1, 1]]> : tensor<2x2xi32> |
| %2 = constant dense<[[[10]], [[20]]]> : tensor<2x1x1xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| "mhlo.return"(%arg1) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 0 : i64, |
| inserted_window_dims = dense<> : tensor<0xi64>, |
| scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>, |
| update_window_dims = dense<[1, 2]> : tensor<2xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2x2xi32>, tensor<2x1x1xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: mhlo.constant dense<[ |
| // CHECK-SAME: [1, 2, 3], [4, 20, 6], [7, 10, 9] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| } |
| |
| // CHECK-LABEL: @scatter_no_update_window_dim |
| func @scatter_no_update_window_dim() -> tensor<3xi32> { |
| %0 = constant dense<[0, 1, 2]> : tensor<3xi32> |
| %1 = constant dense<[[[0], [1]], [[2], [1]]]> : tensor<2x2x1xi32> |
| %2 = constant dense<[[10, 20], [30, 40]]> : tensor<2x2xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| %4 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> (tensor<i32>) |
| "mhlo.return"(%4) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 2 : i64, |
| inserted_window_dims = dense<0> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, |
| update_window_dims = dense<> : tensor<0xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3xi32>, tensor<2x2x1xi32>, tensor<2x2xi32>) -> tensor<3xi32> |
| return %3 : tensor<3xi32> |
| // CHECK: mhlo.constant dense<[10, 61, 32]> : tensor<3xi32> |
| } |
| |
| // CHECK-LABEL: @scatter_negative_index |
| func @scatter_negative_index() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[0, -1]> : tensor<2xi32> |
| %2 = constant dense<[[10, 20, 30], [70, 80, 90]]> : tensor<2x3xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| "mhlo.return"(%arg1) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 1 : i64, |
| inserted_window_dims = dense<0> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2xi32>, tensor<2x3xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: constant dense<[ |
| // CHECK-SAME: [1, 2, 3], [4, 5, 6], [7, 8, 9] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| // CHECK: "mhlo.scatter" |
| } |
| |
| // CHECK-LABEL: @scatter_out_of_bound |
| func @scatter_out_of_bound() -> tensor<3x3xi32> { |
| %0 = constant dense<[[1, 2, 3], [4, 5, 6], [7, 8, 9]]> : tensor<3x3xi32> |
| %1 = constant dense<[1, 5]> : tensor<2xi32> |
| %2 = constant dense<[[10, 20, 30], [70, 80, 90]]> : tensor<2x3xi32> |
| %3 = "mhlo.scatter"(%0, %1, %2) ( { |
| ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>): |
| "mhlo.return"(%arg1) : (tensor<i32>) -> () |
| }) {indices_are_sorted = false, |
| scatter_dimension_numbers = { |
| index_vector_dim = 1 : i64, |
| inserted_window_dims = dense<0> : tensor<1xi64>, |
| scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, |
| update_window_dims = dense<[1]> : tensor<1xi64> |
| }, |
| unique_indices = false |
| } : (tensor<3x3xi32>, tensor<2xi32>, tensor<2x3xi32>) -> tensor<3x3xi32> |
| return %3 : tensor<3x3xi32> |
| // CHECK: constant dense<[ |
| // CHECK-SAME: [1, 2, 3], [4, 5, 6], [7, 8, 9] |
| // CHECK-SAME: ]> : tensor<3x3xi32> |
| // CHECK: "mhlo.scatter" |
| } |
| |
| // CHECK-LABEL: @pad_identity_fold |
| func @pad_identity_fold(%arg0: tensor<5x7xf32>) -> tensor<5x7xf32> { |
| %0 = constant dense<0.0> : tensor<f32> |
| %1 = "mhlo.pad"(%arg0, %0) { |
| edge_padding_low = dense<0> : tensor<2xi64>, |
| edge_padding_high = dense<0> : tensor<2xi64>, |
| interior_padding = dense<0> : tensor<2xi64> |
| } : (tensor<5x7xf32>, tensor<f32>) -> tensor<5x7xf32> |
| return %1 : tensor<5x7xf32> |
| // CHECK: return %arg0 : tensor<5x7xf32> |
| } |
| |
| // CHECK-LABEL: @pad_fold |
| func @pad_fold() -> tensor<4x5xi32> { |
| %0 = constant dense<[[2, 3], [4, 5]]> : tensor<2x2xi32> |
| %1 = constant dense<1> : tensor<i32> |
| %3 = "mhlo.pad"(%0, %1) { |
| edge_padding_low = dense<[1, 0]> : tensor<2xi64>, |
| edge_padding_high = dense<[1, 2]> : tensor<2xi64>, |
| interior_padding = dense<[0, 1]> : tensor<2xi64> |
| } : (tensor<2x2xi32>, tensor<i32>) -> tensor<4x5xi32> |
| return %3 : tensor<4x5xi32> |
| // CHECK: constant dense<[ |
| // CHECK-SAME: [1, 1, 1, 1, 1], [2, 1, 3, 1, 1], [4, 1, 5, 1, 1], [1, 1, 1, 1, 1] |
| // CHECK-SAME: ]> : tensor<4x5xi32> |
| } |
| |
| func @pad_fold_zero_elements() -> tensor<3xi32> { |
| %0 = mhlo.constant dense<> : tensor<0xi32> |
| %1 = mhlo.constant dense<7> : tensor<i32> |
| %2 = "mhlo.pad"(%0, %1) {edge_padding_high = dense<3> : tensor<1xi64>, edge_padding_low = dense<0> : tensor<1xi64>, interior_padding = dense<0> : tensor<1xi64>} : (tensor<0xi32>, tensor<i32>) -> tensor<3xi32> |
| return %2 : tensor<3xi32> |
| // CHECK: mhlo.constant dense<7> : tensor<3xi32> |
| } |
| |
| // CHECK-LABEL: @identity_broadcast_reshape |
| func @identity_broadcast_reshape(%arg0: tensor<128xf32>) -> tensor<128xf32> { |
| %0 = "mhlo.broadcast"(%arg0) { |
| broadcast_sizes = dense<[1]> : tensor<1xi64>} : (tensor<128xf32>) -> tensor<1x128xf32> |
| %1 = "mhlo.reshape"(%0) : (tensor<1x128xf32>) -> tensor<128xf32> |
| return %1 : tensor<128xf32> |
| // CHECK: return %arg0 : tensor<128xf32> |
| } |
| |
| // CHECK-LABEL: @identity_broadcast_in_dim_reshape |
| func @identity_broadcast_in_dim_reshape(%arg0: tensor<128xf32>) -> tensor<128xf32> { |
| %0 = "mhlo.broadcast_in_dim"(%arg0) { |
| broadcast_dimensions = dense<[1]> : tensor<1xi64> } : (tensor<128xf32>) -> tensor<1x128xf32> |
| %1 = "mhlo.reshape"(%0) : (tensor<1x128xf32>) -> tensor<128xf32> |
| return %1 : tensor<128xf32> |
| // CHECK: return %arg0 : tensor<128xf32> |
| } |
| |
| // CHECK-LABEL: @broadcast_of_reshape |
| func @broadcast_of_reshape(%arg: tensor<?xf32>, |
| %shape: tensor<2xindex>) -> tensor<?x?xf32> { |
| %0 = "mhlo.dynamic_reshape"(%arg, %shape) |
| : (tensor<?xf32>, tensor<2xindex>) -> tensor<?x?xf32> |
| %1 = "mhlo.dynamic_broadcast_in_dim"(%0, %shape) { |
| broadcast_dimensions = dense<[0, 1]> : tensor<2xi64> |
| } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32> |
| return %1 : tensor<?x?xf32> |
| } |
| // CHECK: [[RESHAPE:%.*]] = "mhlo.dynamic_reshape" |
| // CHECK: return [[RESHAPE]] |
| |
| // CHECK-LABEL: @permutation_broadcast_of_reshape |
| func @permutation_broadcast_of_reshape(%arg: tensor<?xf32>, |
| %shape: tensor<2xindex>) -> tensor<?x?xf32> { |
| %0 = "mhlo.dynamic_reshape"(%arg, %shape) |
| : (tensor<?xf32>, tensor<2xindex>) -> tensor<?x?xf32> |
| %1 = "mhlo.dynamic_broadcast_in_dim"(%0, %shape) { |
| broadcast_dimensions = dense<[1, 0]> : tensor<2xi64> |
| } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32> |
| return %1 : tensor<?x?xf32> |
| } |
| // CHECK: mhlo.dynamic_reshape |
| // CHECK: mhlo.dynamic_broadcast_in_dim |
| |
| // CHECK-LABEL: @reshape_of_same_shape_op_result |
| func @reshape_of_same_shape_op_result(%arg: tensor<?xf32>, |
| %shape: tensor<2xindex>) -> tensor<?x?xf32> { |
| %0 = "mhlo.dynamic_reshape"(%arg, %shape) |
| : (tensor<?xf32>, tensor<2xindex>) -> tensor<?x?xf32> |
| %1 = "mhlo.abs"(%0) : (tensor<?x?xf32>) -> tensor<?x?xf32> |
| %2 = "mhlo.dynamic_reshape"(%1, %shape) |
| : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32> |
| return %2 : tensor<?x?xf32> |
| } |
| // CHECK: mhlo.dynamic_reshape |
| // CHECK-NEXT: mhlo.abs |
| // CHECK-NOT: mhlo.dynamic_reshape |
