1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' -verify-diagnostics \ 2// RUN: -split-input-file | FileCheck %s 3 4func.func @unknown_op() -> index { 5 %0 = "test.foo"() : () -> (tensor<?x?xf32>) 6 // expected-error @below{{could not reify bound}} 7 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?x?xf32>) -> (index) 8 return %1 : index 9} 10 11// ----- 12 13// CHECK-LABEL: func @cast( 14// CHECK: %[[c10:.*]] = arith.constant 10 : index 15// CHECK: return %[[c10]] 16func.func @cast(%t: tensor<10xf32>) -> index { 17 %0 = tensor.cast %t : tensor<10xf32> to tensor<?xf32> 18 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index) 19 return %1 : index 20} 21 22// ----- 23 24func.func @cast_unranked(%t: tensor<*xf32>) -> index { 25 %0 = tensor.cast %t : tensor<*xf32> to tensor<?xf32> 26 // expected-error @below{{could not reify bound}} 27 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index) 28 return %1 : index 29} 30 31// ----- 32 33// CHECK-LABEL: func @dim( 34// CHECK-SAME: %[[t:.*]]: tensor<?xf32> 35// CHECK: %[[dim:.*]] = tensor.dim %[[t]] 36// CHECK: %[[dim:.*]] = tensor.dim %[[t]] 37// CHECK: return %[[dim]] 38func.func @dim(%t: tensor<?xf32>) -> index { 39 %c0 = arith.constant 0 : index 40 %0 = tensor.dim %t, %c0 : tensor<?xf32> 41 %1 = "test.reify_bound"(%0) : (index) -> (index) 42 return %1 : index 43} 44 45// ----- 46 47// CHECK-LABEL: func @dim_all_positive( 48func.func @dim_all_positive(%t: tensor<?xf32>, %x: index) { 49 %c0 = arith.constant 0 : index 50 %0 = tensor.dim %t, %x : tensor<?xf32> 51 // expected-remark @below{{true}} 52 "test.compare"(%0, %c0) {cmp = "GE" } : (index, index) -> () 53 return 54} 55 56// ----- 57 58// CHECK-LABEL: func @empty( 59// CHECK-SAME: %[[sz:.*]]: index 60// CHECK: %[[c6:.*]] = arith.constant 6 : index 61// CHECK: return %[[c6]], %[[sz]] 62func.func @empty(%sz: index) -> (index, index) { 63 %0 = tensor.empty(%sz) : tensor<6x?xf32> 64 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<6x?xf32>) -> (index) 65 %2 = "test.reify_bound"(%0) {dim = 1} : (tensor<6x?xf32>) -> (index) 66 return %1, %2 : index, index 67} 68 69// ----- 70 71// CHECK-LABEL: func @extract_slice_dynamic( 72// CHECK-SAME: %[[t:.*]]: tensor<?xf32>, %[[sz:.*]]: index 73// CHECK: return %[[sz]] 74func.func @extract_slice_dynamic(%t: tensor<?xf32>, %sz: index) -> index { 75 %0 = tensor.extract_slice %t[2][%sz][1] : tensor<?xf32> to tensor<?xf32> 76 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index) 77 return %1 : index 78} 79 80// ----- 81 82// CHECK-LABEL: func @extract_slice_static( 83// CHECK-SAME: %[[t:.*]]: tensor<?xf32> 84// CHECK: %[[c5:.*]] = arith.constant 5 : index 85// CHECK: return %[[c5]] 86func.func @extract_slice_static(%t: tensor<?xf32>) -> index { 87 %0 = tensor.extract_slice %t[2][5][1] : tensor<?xf32> to tensor<5xf32> 88 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<5xf32>) -> (index) 89 return %1 : index 90} 91 92// ----- 93 94func.func @extract_slice_dynamic_constant(%t: tensor<?xf32>, %sz: index) -> index { 95 %0 = tensor.extract_slice %t[2][%sz][1] : tensor<?xf32> to tensor<?xf32> 96 // expected-error @below{{could not reify bound}} 97 %1 = "test.reify_bound"(%0) {dim = 0, constant} : (tensor<?xf32>) -> (index) 98 return %1 : index 99} 100 101// ----- 102 103// CHECK-LABEL: func @extract_slice_static_constant( 104// CHECK-SAME: %[[t:.*]]: tensor<?xf32> 105// CHECK: %[[c5:.*]] = arith.constant 5 : index 106// CHECK: return %[[c5]] 107func.func @extract_slice_static_constant(%t: tensor<?xf32>) -> index { 108 %0 = tensor.extract_slice %t[2][5][1] : tensor<?xf32> to tensor<5xf32> 109 %1 = "test.reify_bound"(%0) {dim = 0, constant} : (tensor<5xf32>) -> (index) 110 return %1 : index 111} 112 113// ----- 114 115// CHECK-LABEL: func @extract_slice_rank_reduce( 116// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>, %[[sz:.*]]: index 117// CHECK: return %[[sz]] 118func.func @extract_slice_rank_reduce(%t: tensor<?x?xf32>, %sz: index) -> index { 119 %0 = tensor.extract_slice %t[0, 2][1, %sz][1, 1] : tensor<?x?xf32> to tensor<?xf32> 120 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index) 121 return %1 : index 122} 123 124// ----- 125 126// CHECK-LABEL: func @insert( 127// CHECK-SAME: %[[t:.*]]: tensor<?xf32> 128// CHECK: %[[c0:.*]] = arith.constant 0 : index 129// CHECK: %[[dim:.*]] = tensor.dim %[[t]], %[[c0]] 130// CHECK: return %[[dim]] 131func.func @insert(%t: tensor<?xf32>, %f: f32, %pos: index) -> index { 132 %0 = tensor.insert %f into %t[%pos] : tensor<?xf32> 133 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index) 134 return %1 : index 135} 136 137// ----- 138 139// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)> 140// CHECK: #[[$map1:.*]] = affine_map<()[s0] -> (s0 + 12)> 141// CHECK-LABEL: func @pad( 142// CHECK-SAME: %[[t:.*]]: tensor<?x7xf32>, %[[a:.*]]: index, %[[b:.*]]: index 143// CHECK: %[[c0:.*]] = arith.constant 0 : index 144// CHECK: %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]] 145// CHECK: %[[bound0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[a]]] 146// CHECK: %[[bound1:.*]] = affine.apply #[[$map1]]()[%[[b]]] 147// CHECK: return %[[bound0]], %[[bound1]] 148func.func @pad(%t: tensor<?x7xf32>, %a: index, %b: index) -> (index, index) { 149 %pad = arith.constant 0.0 : f32 150 %0 = tensor.pad %t low[%a, 5] high[%a, %b] { 151 ^bb0(%arg1: index, %arg2: index): 152 tensor.yield %pad : f32 153 } : tensor<?x7xf32> to tensor<?x?xf32> 154 %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?x?xf32>) -> (index) 155 %2 = "test.reify_bound"(%0) {dim = 1} : (tensor<?x?xf32>) -> (index) 156 return %1, %2 : index, index 157} 158 159// ----- 160 161// CHECK-LABEL: func @rank( 162// CHECK-SAME: %[[t:.*]]: tensor<5xf32> 163// CHECK: %[[c1:.*]] = arith.constant 1 : index 164// CHECK: return %[[c1]] 165func.func @rank(%t: tensor<5xf32>) -> index { 166 %0 = tensor.rank %t : tensor<5xf32> 167 %1 = "test.reify_bound"(%0) : (index) -> (index) 168 return %1 : index 169} 170 171// ----- 172 173func.func @dynamic_dims_are_equal(%t: tensor<?xf32>) { 174 %c0 = arith.constant 0 : index 175 %dim0 = tensor.dim %t, %c0 : tensor<?xf32> 176 %dim1 = tensor.dim %t, %c0 : tensor<?xf32> 177 // expected-remark @below {{true}} 178 "test.compare"(%dim0, %dim1) : (index, index) -> () 179 return 180} 181 182// ----- 183 184func.func @dynamic_dims_are_different(%t: tensor<?xf32>) { 185 %c0 = arith.constant 0 : index 186 %c1 = arith.constant 1 : index 187 %dim0 = tensor.dim %t, %c0 : tensor<?xf32> 188 %val = arith.addi %dim0, %c1 : index 189 // expected-remark @below {{false}} 190 "test.compare"(%dim0, %val) : (index, index) -> () 191 return 192} 193 194// ----- 195 196func.func @dynamic_dims_are_maybe_equal_1(%t: tensor<?xf32>) { 197 %c0 = arith.constant 0 : index 198 %c5 = arith.constant 5 : index 199 %dim0 = tensor.dim %t, %c0 : tensor<?xf32> 200 // expected-error @below {{unknown}} 201 "test.compare"(%dim0, %c5) : (index, index) -> () 202 return 203} 204 205// ----- 206 207func.func @dynamic_dims_are_maybe_equal_2(%t: tensor<?x?xf32>) { 208 %c0 = arith.constant 0 : index 209 %c1 = arith.constant 1 : index 210 %dim0 = tensor.dim %t, %c0 : tensor<?x?xf32> 211 %dim1 = tensor.dim %t, %c1 : tensor<?x?xf32> 212 // expected-error @below {{unknown}} 213 "test.compare"(%dim0, %dim1) : (index, index) -> () 214 return 215} 216